Initial generation of IR from nu AST. (#40)

The initial implementation only handles simple math expressions, the
focus is primarily on introducing new structs, agreeing on naming and
adding the output to insta tests.

Author: doriath

Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 [dependencies]
 tracy-client = { version = "0.17.3", default-features = false } # for tracy v0.11.1
 logos = "0.15"
+nu-protocol = "0.101"
 
 [profile.profiling]
 inherits = "release"
@@ -30,7 +31,6 @@ path = "src/lib.rs"
 insta = { version = "1.33.0", features = ["glob"] }
 tango-bench = "0.6"
 nu-parser = "0.101"
-nu-protocol = "0.101"
 nu-cmd-lang = "0.101"
 
 [[bench]]

src/compiler.rs

@@ -233,4 +233,17 @@ impl Compiler {
             .get(span_start..span_end)
             .expect("internal error: missing source of span")
     }
+
+    /// Get the source contents of a node
+    pub fn node_as_str(&self, node_id: NodeId) -> &str {
+        std::str::from_utf8(self.get_span_contents(node_id))
+            .expect("internal error: expected utf8 string")
+    }
+
+    /// Get the source contents of a node as i64
+    pub fn node_as_i64(&self, node_id: NodeId) -> i64 {
+        self.node_as_str(node_id)
+            .parse::<i64>()
+            .expect("internal error: expected i64")
+    }
 }

src/ir_generator.rs

@@ -0,0 +1,170 @@
+use crate::compiler::Compiler;
+use crate::errors::{Severity, SourceError};
+use crate::parser::{AstNode, NodeId};
+use nu_protocol::ast::{Math, Operator};
+use nu_protocol::ir::{Instruction, IrBlock, Literal};
+use nu_protocol::{RegId, Span};
+
+/// Generates IR (Intermediate Representation) from nu AST.
+pub struct IrGenerator<'a> {
+    // Immutable reference to a compiler after the typechecker pass
+    compiler: &'a Compiler,
+    errors: Vec<SourceError>,
+    block: IrBlock,
+}
+
+impl<'a> IrGenerator<'a> {
+    pub fn new(compiler: &'a Compiler) -> Self {
+        Self {
+            compiler,
+            errors: Default::default(),
+            block: IrBlock {
+                instructions: Default::default(),
+                spans: Default::default(),
+                data: Default::default(),
+                ast: Default::default(),
+                comments: Default::default(),
+                register_count: 0,
+                file_count: 0,
+            },
+        }
+    }
+
+    /// Generates the IR from the given state of the compiler.
+    /// After this is called, use `block` and `errors` to get the result.
+    pub fn generate(&mut self) {
+        if self.compiler.ast_nodes.is_empty() {
+            return;
+        }
+        let node_id = NodeId(self.compiler.ast_nodes.len() - 1);
+        let Some(reg) = self.generate_node(node_id) else {
+            return;
+        };
+        self.add_instruction(node_id, Instruction::Return { src: reg });
+    }
+
+    /// Returns generated IR block.
+    ///
+    /// Call `generate` before using this method and ensure there are no errors.
+    pub fn block(self) -> IrBlock {
+        self.block
+    }
+
+    /// Returns errors encountered during IR generation step.
+    ///
+    /// Call `generate` before using this method.
+    pub fn errors(&self) -> &Vec<SourceError> {
+        &self.errors
+    }
+
+    /// Prints the internal state to standard output.
+    pub fn print(&self) {
+        let output = self.display_state();
+        print!("{output}");
+    }
+
+    /// Displays the state of the IR generator.
+    /// The output can be used for human debugging and for snapshot tests.
+    pub fn display_state(&self) -> String {
+        let mut result = String::new();
+        result.push_str("==== IR ====\n");
+        result.push_str(&format!("register_count: {}\n", self.block.register_count));
+        result.push_str(&format!("file_count: {}\n", self.block.file_count));
+
+        for (idx, instruction) in self.block.instructions.iter().enumerate() {
+            result.push_str(&format!("{}: {:?}\n", idx, instruction));
+        }
+
+        if !self.errors.is_empty() {
+            result.push_str("==== IR ERRORS ====\n");
+            for error in &self.errors {
+                result.push_str(&format!(
+                    "{:?} (NodeId {}): {}\n",
+                    error.severity, error.node_id.0, error.message
+                ));
+            }
+        }
+        result
+    }
+
+    // Returns unused register.
+    fn next_register(&mut self) -> RegId {
+        let r = RegId::new(self.block.register_count);
+        self.block.register_count += 1;
+        r
+    }
+
+    fn generate_node(&mut self, node_id: NodeId) -> Option<RegId> {
+        let ast_node = &self.compiler.ast_nodes[node_id.0];
+        match ast_node {
+            AstNode::Int => {
+                let next_reg = self.next_register();
+                let val = self.compiler.node_as_i64(node_id);
+                self.add_instruction(
+                    node_id,
+                    Instruction::LoadLiteral {
+                        dst: next_reg,
+                        lit: Literal::Int(val),
+                    },
+                );
+                Some(next_reg)
+            }
+            AstNode::Block(block_id) => {
+                let block = &self.compiler.blocks[block_id.0];
+                let mut last = None;
+                for id in &block.nodes {
+                    last = self.generate_node(*id);
+                    last?;
+                }
+                last
+            }
+            AstNode::BinaryOp { lhs, op, rhs } => {
+                let l = self.generate_node(*lhs)?;
+                let r = self.generate_node(*rhs)?;
+                let o = self.node_to_operator(*op)?;
+                self.add_instruction(
+                    node_id,
+                    Instruction::BinaryOp {
+                        lhs_dst: l,
+                        op: o,
+                        rhs: r,
+                    },
+                );
+                Some(l)
+            }
+            _ => {
+                self.error(format!("node {:?} not suported yet", ast_node), node_id);
+                None
+            }
+        }
+    }
+
+    fn add_instruction(&mut self, node_id: NodeId, instruction: Instruction) {
+        let span = self.compiler.get_span(node_id);
+        self.block.spans.push(Span {
+            start: span.start,
+            end: span.end,
+        });
+        self.block.ast.push(None);
+        self.block.instructions.push(instruction);
+    }
+
+    fn node_to_operator(&mut self, node_id: NodeId) -> Option<Operator> {
+        match self.compiler.get_node(node_id) {
+            AstNode::Plus => Some(Operator::Math(Math::Plus)),
+            AstNode::Multiply => Some(Operator::Math(Math::Multiply)),
+            node => {
+                self.error(format!("unrecognized operator {:?}", node), node_id);
+                None
+            }
+        }
+    }
+
+    fn error(&mut self, message: impl Into<String>, node_id: NodeId) {
+        self.errors.push(SourceError {
+            message: message.into(),
+            node_id,
+            severity: Severity::Error,
+        });
+    }
+}

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod compiler;
 pub mod errors;
+pub mod ir_generator;
 pub mod lexer;
 pub mod parser;
 pub mod protocol;

src/main.rs

@@ -1,6 +1,7 @@
 use std::process::exit;
 
 use new_nu_parser::compiler::Compiler;
+use new_nu_parser::ir_generator::IrGenerator;
 use new_nu_parser::lexer::lex;
 use new_nu_parser::parser::Parser;
 use new_nu_parser::resolver::Resolver;
@@ -80,6 +81,12 @@ fn main() {
             typechecker.print();
         }
 
+        let mut ir_generator = IrGenerator::new(&compiler);
+        ir_generator.generate();
+        if do_print {
+            ir_generator.print();
+        }
+
         compiler.merge_types(typechecker.to_types());
     }
 }

src/snapshots/[email protected]

@@ -23,3 +23,8 @@ snapshot_kind: text
 4: string
 5: stream<binary>
 6: stream<binary>
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Alias { new_name: NodeId(0), old_name: NodeId(1) } not suported yet

src/snapshots/new_nu_parser__test__node_output@binary_ops_exact.nu.snap

@@ -2,6 +2,7 @@
 source: src/test.rs
 expression: evaluate_example(path)
 input_file: tests/binary_ops_exact.nu
+snapshot_kind: text
 ---
 ==== COMPILER ====
 0: Int (0 to 1) "1"
@@ -71,3 +72,10 @@ input_file: tests/binary_ops_exact.nu
 29: list<int>
 30: bool
 31: bool
+==== IR ====
+register_count: 2
+file_count: 0
+0: LoadLiteral { dst: RegId(0), lit: Int(1) }
+1: LoadLiteral { dst: RegId(1), lit: Int(1) }
+==== IR ERRORS ====
+Error (NodeId 1): unrecognized operator Equal

src/snapshots/new_nu_parser__test__node_output@binary_ops_mismatch.nu.snap

@@ -46,3 +46,8 @@ Error (NodeId 1): type mismatch: unsupported addition between string and float
 Error (NodeId 5): type mismatch: unsupported append between string and float
 Error (NodeId 9): type mismatch: unsupported logical operation between bool and string
 Error (NodeId 13): type mismatch: unsupported string operation between bool and string
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 0): node String not suported yet

src/snapshots/new_nu_parser__test__node_output@binary_ops_subtypes.nu.snap

@@ -2,6 +2,7 @@
 source: src/test.rs
 expression: evaluate_example(path)
 input_file: tests/binary_ops_subtypes.nu
+snapshot_kind: text
 ---
 ==== COMPILER ====
 0: Int (0 to 1) "1"
@@ -115,3 +116,9 @@ input_file: tests/binary_ops_subtypes.nu
 51: int
 52: float
 53: float
+==== IR ====
+register_count: 1
+file_count: 0
+0: LoadLiteral { dst: RegId(0), lit: Int(1) }
+==== IR ERRORS ====
+Error (NodeId 2): node Float not suported yet

src/snapshots/[email protected]

@@ -90,3 +90,8 @@ snapshot_kind: text
 36: unknown
 37: stream<binary>
 38: stream<binary>
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 7): node Call { parts: [NodeId(0), NodeId(1), NodeId(2), NodeId(6)] } not suported yet