Adds graphviz support to lambo.

3 files changed, 114 insertions, 43 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 762f3ae..b32d733 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -205,6 +205,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "bstr"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c48f0051a4b4c5e0b6d365cd04af53aeaa209e3cc15ec2cdb69e73cc87fbd0dc"
+dependencies = [
+ "memchr",
+ "regex-automata",
+ "serde",
+]
+
+[[package]]
 name = "bumpalo"
 version = "3.14.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1009,6 +1020,7 @@ name = "lambo"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "bstr",
  "clap",
  "csv",
  "futures",
@@ -1423,6 +1435,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "regex-automata"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f"
+
+[[package]]
 name = "reqwest"
 version = "0.11.23"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 72163da..0778143 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 anyhow = { version = "1.0.79", features = ["backtrace"] }
+bstr = "1.9.0"
 clap = { version = "4.4.17", features = ["derive"] }
 csv = "1.3.0"
 futures = "0.3.30"
diff --git a/src/commands/discocaml.rs b/src/commands/discocaml.rs
index 8565c6f..fb6b05b 100644
--- a/src/commands/discocaml.rs
+++ b/src/commands/discocaml.rs
@@ -1,4 +1,5 @@
 use anyhow::{anyhow, bail, Context as _, Error, Result};
+use bstr::BString;
 use futures::Future;
 use serde::{Deserialize, Serialize};
 use serde_json::de::Deserializer;
@@ -9,8 +10,8 @@ use serenity::{
         InteractionResponseFlags, Member, RoleId,
     },
     builder::{
-        CreateButton, CreateCommand, CreateCommandOption, CreateInteractionResponse,
-        CreateInteractionResponseMessage,
+        CreateAttachment, CreateButton, CreateCommand, CreateCommandOption, CreateEmbed,
+        CreateInteractionResponse, CreateInteractionResponseMessage,
     },
     client::Context,
     model::Permissions,
@@ -74,6 +75,7 @@ pub enum DiscocamlCommand {
 pub enum DiscocamlResponse {
     Expr(DiscocamlResponseExpr),
     Error(String),
+    Graphviz(String),
 }
 
 #[derive(Debug, Deserialize, PartialEq)]
@@ -139,10 +141,60 @@ async fn respond_with_error<E: std::error::Error, F: Future<Output = Result<(),
     }
 }
 
+async fn respond_with<E, F>(
+    db: &SqlitePool,
+    interaction_id: InteractionId,
+    res: DiscocamlResponse,
+    send: impl FnOnce(CreateInteractionResponse) -> F,
+) -> Result<()>
+where
+    E: std::error::Error + 'static + Send + Sync,
+    F: Future<Output = Result<(), E>>,
+{
+    match res {
+        DiscocamlResponse::Expr(expr) => {
+            // Insert the output expression in the database.
+            if let Err(err) = expr
+                .save(db, interaction_id)
+                .await
+                .context("failed to save expression to database")
+            {
+                respond_with_error(&err, send).await;
+                return Err(err);
+            }
+
+            // Respond with the expression and the buttons.
+            let res = expr_response_message(&expr);
+            send(CreateInteractionResponse::Message(res)).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: `{}`", err);
+            respond_with_error(&err, send).await;
+            return Err(err);
+        }
+    }
+}
+
 pub async fn run_discocaml(
     config: &DiscocamlConfig,
     req: &DiscocamlRequest,
-) -> Result<DiscocamlResponseExpr> {
+) -> Result<DiscocamlResponse> {
     let mut child = Command::new(&config.command[0])
         .args(&config.command)
         .stdin(Stdio::piped())
@@ -177,14 +229,40 @@ pub async fn run_discocaml(
         .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<BString> {
+    let mut child = Command::new("dot")
+        .arg("-Tpng")
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .kill_on_drop(true)
+        .spawn()
+        .context("failed to start dot")?;
 
-    match out {
-        DiscocamlResponse::Expr(expr) => Ok(expr),
-        DiscocamlResponse::Error(err) => bail!("got an error from discocaml: {:?}", err),
+    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 make_response_message(expr: &DiscocamlResponseExpr) -> CreateInteractionResponseMessage {
+fn expr_response_message(expr: &DiscocamlResponseExpr) -> CreateInteractionResponseMessage {
     // TODO: Real escaping
     CreateInteractionResponseMessage::new()
         .content(format!("```ocaml\n{}\n```", expr.expr))
@@ -256,24 +334,11 @@ pub async fn handle_command(
         }
     };
 
-    // Insert the output expression in the database.
-    if let Err(err) = res
-        .save(db, command.id)
-        .await
-        .context("failed to save expression to database")
-    {
-        respond_with_error(&err, |res| command.create_response(&ctx, res)).await;
-        return Err(err);
-    }
-
-    // Respond with the expression and the buttons.
-    command
-        .create_response(
-            &ctx,
-            CreateInteractionResponse::Message(make_response_message(&res)),
-        )
-        .await
-        .context("failed to respond")
+    // Respond with the resulting expression.
+    respond_with(db, command.id, res, |res| {
+        command.create_response(&ctx, res)
+    })
+    .await
 }
 
 pub async fn handle_button(
@@ -340,22 +405,9 @@ pub async fn handle_button(
         }
     };
 
-    // Insert the output expression in the database.
-    if let Err(err) = res
-        .save(db, component.id)
-        .await
-        .context("failed to save expression to database")
-    {
-        respond_with_error(&err, |res| component.create_response(&ctx, res)).await;
-        return Err(err);
-    }
-
-    // Respond with the expression and the buttons.
-    component
-        .create_response(
-            &ctx,
-            CreateInteractionResponse::Message(make_response_message(&res)),
-        )
-        .await
-        .context("failed to respond")
+    // Respond.
+    respond_with(db, component.id, res, |res| {
+        component.create_response(&ctx, res)
+    })
+    .await
 }