use anyhow::{anyhow, bail, Context as _, Error, Result}; use bstr::BString; use futures::Future; use serde::{Deserialize, Serialize}; use serde_json::de::Deserializer; use serenity::{ all::{ CommandDataOptionValue, CommandInteraction, CommandOptionType, CommandType, ComponentInteraction, ComponentInteractionDataKind, InteractionResponseFlags, Member, RoleId, }, builder::{ CreateAttachment, CreateButton, CreateCommand, CreateCommandOption, CreateEmbed, CreateInteractionResponse, CreateInteractionResponseMessage, }, client::Context, model::Permissions, }; use sqlx::SqlitePool; use std::{collections::BTreeSet, process::Stdio}; use tokio::{io::AsyncWriteExt, process::Command}; use crate::utils::EnumAsArray; #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct DiscocamlConfig { pub command: Vec, pub roles: BTreeSet, } #[derive(Debug, PartialEq, Serialize)] pub struct DiscocamlRequest { pub expr: String, pub command: DiscocamlCommand, } #[derive(Debug, PartialEq, Serialize)] pub enum DiscocamlCommand { Parse, DrawTree, RunCBN, RunCBV, StepCBN, StepCBV, } /// A response outputted by the discocaml subprocess as a JSON string. /// /// ``` /// # use lambo::{commands::discocaml::*, utils::EnumAsArray}; /// # use serde::Deserialize; /// # use serde_json::Deserializer; /// /// let example = r#" /// [ "Expr" /// , { "expr": "1 + 2" /// , "has_redex": true /// , "note": null /// } /// ] /// "#; /// let expected = DiscocamlResponse::Expr(DiscocamlResponseExpr { /// expr: "1 + 2".to_string(), /// has_redex: true, /// note: None, /// }); /// /// let mut de = Deserializer::from_str(&example); /// let out = DiscocamlResponse::deserialize(EnumAsArray(&mut de)).unwrap(); /// de.end().unwrap(); /// /// assert_eq!(out, expected); /// ``` /// /// ``` /// # use lambo::{commands::discocaml::*, utils::EnumAsArray}; /// # use serde::Deserialize; /// # use serde_json::Deserializer; /// /// let example = r#"["Exprs", "A", ["B", "C"]]"#; /// let expected = DiscocamlResponse::Exprs( /// "A".to_string(), /// vec![ /// "B".to_string(), /// "C".to_string(), /// ], /// ); /// /// let mut de = Deserializer::from_str(&example); /// let out = DiscocamlResponse::deserialize(EnumAsArray(&mut de)).unwrap(); /// de.end().unwrap(); /// /// assert_eq!(out, expected); /// ``` #[derive(Debug, Deserialize, PartialEq)] #[serde(deny_unknown_fields)] pub enum DiscocamlResponse { Expr(DiscocamlResponseExpr), Exprs(Option, Vec), Error(String), Graphviz(String), } #[derive(Debug, Deserialize, PartialEq)] #[serde(deny_unknown_fields)] pub struct DiscocamlResponseExpr { pub expr: String, pub has_redex: bool, pub note: Option, } impl DiscocamlResponseExpr { async fn save(&self, db: &SqlitePool) -> Result { let out = sqlx::query!( "INSERT INTO discocaml_exprs (expr) VALUES (?) RETURNING rowid", self.expr ) .fetch_one(db) .await .context("failed to save expression to database")?; Ok(out.rowid) } } pub fn command() -> CreateCommand { CreateCommand::new("discocaml") .kind(CommandType::ChatInput) .default_member_permissions(Permissions::empty()) .dm_permission(true) .description("Sends this expression to the disco!") .add_option( CreateCommandOption::new( CommandOptionType::String, "expr", "The expression to operate on.", ) .required(true), ) } async fn check_permissions(config: &DiscocamlConfig, member: Option<&Member>) -> Result<()> { if let Some(member) = member { if !member.roles.iter().any(|role| config.roles.contains(role)) { match config.roles.len() { 0 => bail!("This command can not be used, because roles are not configured."), 1 => bail!( "This command can only be used by <@&{}>.", config.roles.iter().next().unwrap() ), 2 => { let mut iter = config.roles.iter(); bail!( "This command can only be used by <@&{}> or <@&{}>.", iter.next().unwrap(), iter.next().unwrap() ) } _ => { let mut iter = config.roles.iter(); let last = iter.next().unwrap(); let mut buf = String::new(); for role in iter { buf.push_str("<@&"); buf.push_str(&role.to_string()); buf.push_str(">, "); } buf.push_str(", or <@&"); buf.push_str(&last.to_string()); buf.push('>'); bail!("This command can only be used by {}.", buf) } } } Ok(()) } else { bail!("This command cannot be used in DMs.") } } async fn respond_with_error>>( err: &Error, send: impl FnOnce(CreateInteractionResponse) -> F, ) { let mut content = format!(":no_entry_sign: {}", err); if let Some((i, _)) = content.char_indices().nth(1997) { content.truncate(i); assert_eq!(content.chars().count(), 1997); content.push('…'); } let msg = CreateInteractionResponseMessage::new() .content(content) .flags(InteractionResponseFlags::EPHEMERAL); if let Err(err) = send(CreateInteractionResponse::Message(msg)).await { log::error!( "failed to respond to command that failed permissions check: {:?}", err ) } } async fn respond_with( db: &SqlitePool, res: DiscocamlResponse, send: impl FnOnce(CreateInteractionResponse) -> F, ) -> Result<()> where E: std::error::Error + 'static + Send + Sync, F: Future>, { match res { DiscocamlResponse::Expr(expr) => { // Insert the output expression in the database. let rowid = match expr.save(db).await { Ok(rowid) => rowid, Err(err) => { respond_with_error(&err, send).await; return Err(err); } }; // Respond with the expression and the buttons. let res = expr_response_message(&expr, rowid); send(CreateInteractionResponse::Message(res)).await?; Ok(()) } DiscocamlResponse::Exprs(note, exprs) => { let mut msg = String::new(); if let Some(note) = ¬e { msg.push_str(note); msg.push('\n'); } let mut msg_chars = msg.chars().count(); for expr in exprs { let mut bullet = String::new(); bullet.push_str("1. ```ocaml\n"); for line in escape_code(&expr).lines() { bullet.push_str(" "); bullet.push_str(line); bullet.push('\n'); } bullet.push_str(" ```\n"); let bullet_chars = bullet.chars().count(); if msg_chars + bullet_chars > 1995 { msg.push_str("1. …\n"); msg_chars += 5; assert!(msg_chars < 2000); break; } msg.push_str(&bullet); msg_chars += bullet_chars; } assert_eq!(msg.chars().count(), msg_chars); let msg = CreateInteractionResponseMessage::new().content(msg); send(CreateInteractionResponse::Message(msg)).await?; Ok(()) } DiscocamlResponse::Graphviz(dot) => { let png = match run_dot(&dot).await { Ok(png) => png, Err(err) => { respond_with_error(&err, send).await; return Err(err); } }; let msg = CreateInteractionResponseMessage::new() .add_file(CreateAttachment::bytes(png, "expr.png")) .add_embed(CreateEmbed::new().attachment("expr.png")); send(CreateInteractionResponse::Message(msg)).await?; Ok(()) } DiscocamlResponse::Error(err) => { let err = anyhow!("got an error from discocaml: `{}`", escape_code(&err)); respond_with_error(&err, send).await; Err(err) } } } pub async fn run_discocaml( config: &DiscocamlConfig, req: &DiscocamlRequest, ) -> Result { let mut child = Command::new(&config.command[0]) .args(&config.command) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .kill_on_drop(true) .spawn() .context("failed to start discocaml")?; let mut stdin = child.stdin.take().unwrap(); let req = serde_json::to_string(req).context("failed to serialize request to discocaml")?; stdin .write_all(req.as_bytes()) .await .context("failed to write request to discocaml")?; stdin .shutdown() .await .context("failed to close pipe to discocaml")?; drop(stdin); drop(req); let output = child .wait_with_output() .await .context("failed to wait for discocaml to complete")?; if !output.status.success() { bail!("discocaml exited with non-zero status {:?}", output.status) } let mut de = Deserializer::from_slice(&output.stdout); let out = DiscocamlResponse::deserialize(EnumAsArray(&mut de)) .context("failed to parse response from discocaml")?; de.end() .context("failed to parse response from discocaml")?; Ok(out) } async fn run_dot(dot: &str) -> Result { let mut child = Command::new("dot") .arg("-Tpng") .stdin(Stdio::piped()) .stdout(Stdio::piped()) .kill_on_drop(true) .spawn() .context("failed to start dot")?; let mut stdin = child.stdin.take().unwrap(); stdin .write_all(dot.as_bytes()) .await .context("failed to write request to dot")?; stdin .shutdown() .await .context("failed to close pipe to dot")?; drop(stdin); let output = child .wait_with_output() .await .context("failed to wait for dot to complete")?; if !output.status.success() { bail!("dot exited with non-zero status {:?}", output.status) } Ok(BString::new(output.stdout)) } fn escape_code(s: &str) -> String { s.replace('`', "\\`") } fn expr_response_message( expr: &DiscocamlResponseExpr, rowid: i64, ) -> CreateInteractionResponseMessage { let mut msg = String::new(); if let Some(note) = &expr.note { msg.push_str(note); msg.push('\n'); } msg.push_str("```ocaml\n"); msg.push_str(&escape_code(&expr.expr)); msg.push_str("```\n"); CreateInteractionResponseMessage::new() .content(msg) .button(CreateButton::new(format!("discocaml-draw-tree/{}", rowid)).label("Draw Tree")) .button( CreateButton::new(format!("discocaml-step-cbv/{}", rowid)) .label("Step") .disabled(!expr.has_redex), ) .button( CreateButton::new(format!("discocaml-run-cbv/{}", rowid)) .label("Run") .disabled(!expr.has_redex), ) /* .button( CreateButton::new(format!("discocaml-step-cbv/{}", rowid)) .label("Step (CBV)") .disabled(!expr.has_redex), ) .button( CreateButton::new(format!("discocaml-step-cbn/{}", rowid)) .label("Step (CBN)") .disabled(!expr.has_redex), ) .button( CreateButton::new(format!("discocaml-run-cbv/{}", rowid)) .label("Run (CBV)") .disabled(!expr.has_redex), ) .button( CreateButton::new(format!("discocaml-run-cbn/{}", rowid)) .label("Run (CBN)") .disabled(!expr.has_redex), ) */ } pub async fn handle_command( ctx: &Context, config: &DiscocamlConfig, db: &SqlitePool, command: &CommandInteraction, ) -> Result<()> { // Check that the required role was present. if let Err(err) = check_permissions(config, command.member.as_deref()).await { respond_with_error(&err, |res| command.create_response(&ctx, res)).await; return Err(err.context("permissions check failed")); } // Parse the expression out. let mut expr = None; for option in &command.data.options { match (&option.name as &str, &option.value) { ("expr", CommandDataOptionValue::String(value)) => expr = Some(value), _ => { let err = anyhow!("unknown option {:?}", option); respond_with_error(&err, |res| command.create_response(&ctx, res)).await; return Err(err); } } } let expr = if let Some(expr) = expr { expr } else { let err = anyhow!("missing option {:?}", "expr"); respond_with_error(&err, |res| command.create_response(&ctx, res)).await; return Err(err); }; // Round-trip the expression through discocaml. let req = DiscocamlRequest { expr: expr.to_string(), command: DiscocamlCommand::Parse, }; let res = match run_discocaml(config, &req).await { Ok(res) => res, Err(err) => { respond_with_error(&err, |res| command.create_response(&ctx, res)).await; return Err(err.context("failed to run discocaml")); } }; // Respond with the resulting expression. respond_with(db, res, |res| command.create_response(&ctx, res)).await } pub async fn handle_button( ctx: &Context, config: &DiscocamlConfig, db: &SqlitePool, component: &ComponentInteraction, ) -> Result<()> { // Check that the required role was present. if let Err(err) = check_permissions(config, component.member.as_ref()).await { respond_with_error(&err, |res| component.create_response(&ctx, res)).await; return Err(err.context("permissions check failed")); } // Come up with a command for discocaml based on which button was clicked. let (command, id) = match &component.data.kind { ComponentInteractionDataKind::Button => { let custom_id = &component.data.custom_id; if let Some(id_str) = custom_id.strip_prefix("discocaml-draw-tree/") { (DiscocamlCommand::DrawTree, id_str) } else if let Some(id_str) = custom_id.strip_prefix("discocaml-step-cbv/") { (DiscocamlCommand::StepCBV, id_str) } else if let Some(id_str) = custom_id.strip_prefix("discocaml-step-cbn/") { (DiscocamlCommand::StepCBN, id_str) } else if let Some(id_str) = custom_id.strip_prefix("discocaml-run-cbv/") { (DiscocamlCommand::RunCBV, id_str) } else if let Some(id_str) = custom_id.strip_prefix("discocaml-run-cbn/") { (DiscocamlCommand::RunCBN, id_str) } else { let err = anyhow!("unknown component {:?}", component.data); respond_with_error(&err, |res| component.create_response(&ctx, res)).await; return Err(err); } } _ => { let err = anyhow!("unknown component {:?}", component.data); respond_with_error(&err, |res| component.create_response(&ctx, res)).await; return Err(err); } }; // Parse the ID, and fetch the corresponding expr. let id: i64 = id.parse().context("failed to parse id")?; let result = sqlx::query!("SELECT expr FROM discocaml_exprs WHERE rowid = ?", id) .fetch_one(db) .await .context("failed to load expression from database"); let expr = match result { Ok(expr) => expr.expr, Err(err) => { respond_with_error(&err, |res| component.create_response(&ctx, res)).await; return Err(err); } }; // Send discocaml the request. let req = DiscocamlRequest { expr, command }; let res = match run_discocaml(config, &req).await { Ok(res) => res, Err(err) => { respond_with_error(&err, |res| component.create_response(&ctx, res)).await; return Err(err.context("failed to run discocaml")); } }; // Respond. respond_with(db, res, |res| component.create_response(&ctx, res)).await }