Merge Egor's Mermaid (#2381)

<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> This pull request integrates Mermaid diagrams for function flows in
the BAML project, enhancing the UI with new components and controls for
diagram interaction and state management.
> 
>   - **Mermaid Integration**:
>     - Add `mermaid` and `svg-pan-zoom` dependencies to `package.json`.
> - Implement `MermaidGraphView` component in `MermaidGraphView.tsx` for
rendering function flow diagrams.
> - Use `mermaid` to generate and render diagrams, with custom CSS for
styling.
>   - **UI Enhancements**:
> - Add zoom and reset controls to `MermaidGraphView` for better user
interaction.
> - Integrate `MermaidGraphView` into `PromptRenderWrapper` with a new
tab for viewing diagrams.
> - Update `PreviewToolbar` to include a toggle for beta features,
enabling the Mermaid graph view.
>   - **State Management**:
> - Introduce `functionGraphAtom` in `atoms-orch-graph.ts` to manage the
state of the function graph.
> - Use Jotai atoms to handle feature flags and runtime configurations.
>   - **Error Handling**:
> - Enhance error rendering in `EnhancedErrorRenderer.tsx` with custom
error renderers and improved UI feedback.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 3c8cb7f. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

---------

Co-authored-by: hellovai <vbv@boundaryml.com>
Co-authored-by: egol <egoriluk@gmail.com>
Co-authored-by: Samuel Lijin <sam@boundaryml.com>
Co-authored-by: Jesús Lapastora <jesus@boundaryml.com>

Author: 5 people

.cursor/rules/jj-workflow.mdc

@@ -1,9 +1,6 @@
 ---
-description: Guide for using fine-grained JJ commits during development.
-globs:
-alwaysApply: true
+alwaysApply: false
 ---
-
 Whenever doing work:
  - See if you are on a bookmark or a commit with a name or description appropriate to that task
  - Use fine-grained tasks - if you're creating tests for TDD, put that in a separate commit

engine/baml-compiler/src/hir/lowering.rs

@@ -375,6 +375,7 @@ fn lower_stmt(stmt: &ast::Stmt) -> Statement {
             is_mutable,
             expr,
             span,
+            annotations: _,
         }) => {
             let lifted_expr = Expression::from_ast(expr);
 
@@ -397,6 +398,7 @@ fn lower_stmt(stmt: &ast::Stmt) -> Statement {
             iterator,
             body,
             span,
+            annotations: _,
         }) => {
             // Lower for loop to HIR
             let lifted_iterator = Expression::from_ast(iterator);
@@ -410,8 +412,8 @@ fn lower_stmt(stmt: &ast::Stmt) -> Statement {
             }
         }
         ast::Stmt::Expression(expr) => Statement::Expression {
-            expr: Expression::from_ast(expr),
-            span: expr.span().clone(),
+            expr: Expression::from_ast(&expr.expr),
+            span: expr.span.clone(),
         },
         ast::Stmt::Semicolon(expr) => Statement::Semicolon {
             expr: Expression::from_ast(expr),

engine/baml-compiler/src/thir.rs

@@ -705,7 +705,7 @@ impl<T: Clone> Statement<T> {
             Statement::Return { expr, span: _ } => {
                 format!("return {}", expr.dump_str())
             }
-            Statement::Expression { expr, span: _ } => expr.dump_str().to_string(),
+            Statement::Expression { expr, span: _ } => expr.dump_str(),
             Statement::SemicolonExpression { expr, span: _ } => expr.dump_str().to_string(),
             Statement::While {
                 condition,

engine/baml-lib/ast/Cargo.toml

@@ -31,6 +31,14 @@ regex.workspace = true
 pretty_assertions.workspace = true
 unindent = "0.2.3"
 
+[[example]]
+name = "generate_mermaid_ast"
+path = "examples/generate_mermaid_ast.rs"
+
+[[example]]
+name = "generate_mermaid_headers"
+path = "examples/generate_mermaid_headers.rs"
+
 [features]
 debug_parser = []
-# default = ["debug_parser"]
+#default = ["debug_parser"]

engine/baml-lib/ast/examples/generate_mermaid_ast.rs

@@ -0,0 +1,77 @@
+use std::{env, fs, path::Path};
+
+use internal_baml_ast::{parse, MermaidDiagramGenerator};
+use internal_baml_diagnostics::SourceFile;
+
+fn main() {
+    // Get command line arguments
+    let args: Vec<String> = env::args().collect();
+
+    if args.len() < 2 || args.len() > 3 {
+        eprintln!("Usage: {} [--fancy] <path-to-baml-file>", args[0]);
+        eprintln!("Example: {} example.baml", args[0]);
+        eprintln!("Example: {} --fancy example.baml", args[0]);
+        std::process::exit(1);
+    }
+
+    let (baml_file_path, use_fancy) = if args.len() == 3 && args[1] == "--fancy" {
+        (&args[2], true)
+    } else {
+        (&args[1], false)
+    };
+
+    // Check if file exists and has .baml extension
+    let path = Path::new(baml_file_path);
+    if !path.exists() {
+        eprintln!("Error: File '{baml_file_path}' does not exist");
+        std::process::exit(1);
+    }
+
+    if path.extension().and_then(|s| s.to_str()) != Some("baml") {
+        eprintln!("Error: File '{baml_file_path}' does not have a .baml extension");
+        std::process::exit(1);
+    }
+
+    // Read the BAML file
+    let baml_source = match fs::read_to_string(path) {
+        Ok(content) => content,
+        Err(err) => {
+            eprintln!("Error reading file '{baml_file_path}': {err}");
+            std::process::exit(1);
+        }
+    };
+
+    // Create a SourceFile
+    let source = SourceFile::from((path.to_path_buf(), baml_source));
+    let root_path = Path::new(".");
+
+    // Parse the BAML source code
+    match parse(root_path, &source) {
+        Ok((ast, _diagnostics)) => {
+            // Print the AST
+            dbg!(&ast);
+
+            // Generate Mermaid diagram with optional styling
+            let mermaid_diagram =
+                MermaidDiagramGenerator::generate_ast_diagram_with_styling(&ast, use_fancy);
+
+            println!(
+                "Generated Mermaid Diagram for '{}' (styling: {}):",
+                baml_file_path,
+                if use_fancy { "enabled" } else { "disabled" }
+            );
+            println!("{mermaid_diagram}");
+
+            // You can copy the output and paste it into any Mermaid renderer
+            println!("\nTo visualize this diagram:");
+            println!("1. Copy the output above");
+            println!("2. Go to https://mermaid.live/");
+            println!("3. Paste the diagram code");
+            println!("4. View the rendered AST diagram!");
+        }
+        Err(err) => {
+            eprintln!("Failed to parse BAML source: {err:?}");
+            std::process::exit(1);
+        }
+    }
+}

engine/baml-lib/ast/examples/generate_mermaid_headers.rs

@@ -0,0 +1,38 @@
+use std::path::Path;
+
+use internal_baml_ast::{parse, BamlVisDiagramGenerator};
+use internal_baml_diagnostics::SourceFile;
+
+fn main() {
+    let mut args = std::env::args().skip(1);
+    let Some(input_path) = args.next() else {
+        eprintln!("Usage: generate_mermaid_headers <path/to/file.baml>");
+        std::process::exit(2);
+    };
+
+    let path = std::path::PathBuf::from(&input_path);
+    let contents = match std::fs::read_to_string(&path) {
+        Ok(s) => s,
+        Err(e) => {
+            eprintln!("Failed to read {input_path}: {e}");
+            std::process::exit(1);
+        }
+    };
+
+    let source = SourceFile::new_allocated(path.clone(), contents.into());
+    let (ast, diags) = match parse(Path::new("."), &source) {
+        Ok(res) => res,
+        Err(diags) => {
+            eprintln!("Parse errors:\n{}", diags.to_pretty_string());
+            std::process::exit(1);
+        }
+    };
+    if diags.has_errors() {
+        eprintln!("Parse errors:\n{}", diags.to_pretty_string());
+        std::process::exit(1);
+    }
+
+    // Nicely styled header graph
+    let mermaid = BamlVisDiagramGenerator::generate_with_styling(&ast, true);
+    println!("{mermaid}");
+}

engine/baml-lib/ast/src/ast.rs

@@ -21,24 +21,31 @@ mod traits;
 mod type_builder_block;
 mod type_expression_block;
 mod value_expression_block;
+
+pub mod baml_vis;
+pub mod header_collector;
+pub mod mermaid_debug;
 pub use app::App;
 pub use argument::{Argument, ArgumentId, ArgumentsList};
 pub use assignment::Assignment;
 pub use attribute::{Attribute, AttributeContainer, AttributeId};
+pub use baml_vis::BamlVisDiagramGenerator;
 pub use config::ConfigBlockProperty;
 pub use expr::{ExprFn, TopLevelAssignment};
 pub use expression::{
     BinaryOperator, ClassConstructor, ClassConstructorField, Expression, ExpressionBlock,
     RawString, UnaryOperator,
 };
 pub use field::{Field, FieldArity, FieldType};
+pub use header_collector::{HeaderCollector, HeaderIndex, RenderableHeader, ScopeId};
 pub use identifier::{Identifier, RefIdentifier};
 pub use indentation_type::IndentationType;
 pub use internal_baml_diagnostics::Span;
+pub use mermaid_debug::MermaidDiagramGenerator;
 pub use newline_type::NewlineType;
 pub use stmt::{
-    AssertStmt, AssignOp, AssignOpStmt, AssignStmt, CForLoopStmt, ForLoopStmt, LetStmt, ReturnStmt,
-    Stmt, WhileStmt,
+    AssertStmt, AssignOp, AssignOpStmt, AssignStmt, CForLoopStmt, ExprStmt, ForLoopStmt, Header,
+    LetStmt, ReturnStmt, Stmt, WhileStmt,
 };
 pub use template_string::TemplateString;
 pub use top::Top;
@@ -103,6 +110,18 @@ impl From<u32> for TypeExpId {
     }
 }
 
+impl From<u32> for ValExpId {
+    fn from(id: u32) -> Self {
+        ValExpId(id)
+    }
+}
+
+impl From<u32> for ExprFnId {
+    fn from(id: u32) -> Self {
+        ExprFnId(id)
+    }
+}
+
 impl std::ops::Index<TypeExpId> for Ast {
     type Output = TypeExpressionBlock;
 
@@ -259,6 +278,13 @@ impl TopId {
         }
     }
 
+    pub fn as_generator_id(self) -> Option<ValExpId> {
+        match self {
+            TopId::Generator(id) => Some(id),
+            _ => None,
+        }
+    }
+
     pub fn as_template_string_id(self) -> Option<TemplateStringId> {
         match self {
             TopId::TemplateString(id) => Some(id),