Initial generation of IR from nu AST.

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/ir_generator.rs

@@ -0,0 +1,193 @@
+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>,
+
+    instructions: Vec<Instruction>,
+    register_count: u32,
+    file_count: u32,
+}
+
+impl<'a> IrGenerator<'a> {
+    pub fn new(compiler: &'a Compiler) -> Self {
+        Self {
+            compiler,
+            errors: Default::default(),
+            instructions: 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 Some(reg) = self.generate_node(NodeId(self.compiler.ast_nodes.len() - 1)) else {
+            return;
+        };
+        self.instructions.push(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 {
+        // TODO: properly generate the spans
+        // TODO: figure out what to do with AST, as this parser has different
+        // representation of AST than the old parser.
+        let mut spans = vec![];
+        let mut ast = vec![];
+        for _ in 0..(self.instructions.len()) {
+            spans.push(Span { start: 0, end: 0 });
+            ast.push(None);
+        }
+        IrBlock {
+            instructions: self.instructions,
+            spans,
+            data: Default::default(),
+            ast,
+            comments: Default::default(),
+            register_count: self.register_count,
+            file_count: self.file_count,
+        }
+    }
+
+    /// 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.register_count));
+        result.push_str(&format!("file_count: {}\n", self.file_count));
+
+        for (idx, instruction) in self.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.register_count);
+        self.register_count += 1;
+        r
+    }
+
+    fn span_to_string(&mut self, node_id: NodeId) -> Option<String> {
+        match std::str::from_utf8(self.compiler.get_span_contents(node_id)) {
+            Ok(val) => Some(val.to_string()),
+            Err(err) => {
+                self.error(
+                    format!("failed to convert a node to string: {err}"),
+                    node_id,
+                );
+                None
+            }
+        }
+    }
+
+    fn span_to_i64(&mut self, node_id: NodeId) -> Option<i64> {
+        match self.span_to_string(node_id)?.parse::<i64>() {
+            Ok(val) => Some(val),
+            Err(err) => {
+                self.error(
+                    format!("failed to convert a node to string: {err}"),
+                    node_id,
+                );
+                None
+            }
+        }
+    }
+
+    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.span_to_i64(node_id)?;
+                self.instructions.push(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.instructions.push(Instruction::BinaryOp {
+                    lhs_dst: l,
+                    op: o,
+                    rhs: r,
+                });
+                Some(l)
+            }
+            _ => {
+                self.error(format!("node {:?} not suported yet", ast_node), node_id);
+                None
+            }
+        }
+    }
+
+    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;
@@ -74,6 +75,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/new_nu_parser__test__node_output@binary_ops_exact.nu.snap

@@ -52,3 +52,10 @@ snapshot_kind: text
 19: bool
 20: bool
 21: 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

@@ -38,3 +38,8 @@ snapshot_kind: text
 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
+==== 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

@@ -96,3 +96,9 @@ snapshot_kind: text
 41: list<list<float>>
 42: list<list<number>>
 43: list<list<number>>
+==== 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/new_nu_parser__test__node_output@calls.nu.snap

@@ -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

src/snapshots/new_nu_parser__test__node_output@closure3.nu.snap

@@ -59,3 +59,8 @@ snapshot_kind: text
 21: closure
 22: stream<binary>
 23: stream<binary>
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 19): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(18), is_mutable: false } not suported yet

src/snapshots/new_nu_parser__test__node_output@def.nu.snap

@@ -61,3 +61,8 @@ snapshot_kind: text
 22: list<any>
 23: ()
 24: ()
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 23): node Def { name: NodeId(0), params: NodeId(17), return_ty: None, block: NodeId(22) } not suported yet

src/snapshots/new_nu_parser__test__node_output@for.nu.snap

@@ -77,3 +77,8 @@ snapshot_kind: text
 29: ()
 30: ()
 31: ()
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(1), is_mutable: true } not suported yet

src/snapshots/new_nu_parser__test__node_output@for_break_continue.nu.snap

@@ -80,3 +80,8 @@ snapshot_kind: text
 ==== TYPE ERRORS ====
 Error (NodeId 12): unsupported ast node 'Break' in typechecker
 Error (NodeId 19): unsupported ast node 'Continue' in typechecker
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(1), is_mutable: true } not suported yet

src/snapshots/new_nu_parser__test__node_output@if_.nu.snap

@@ -52,3 +52,8 @@ snapshot_kind: text
 17: int
 18: int
 19: int
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(1), is_mutable: false } not suported yet

src/snapshots/new_nu_parser__test__node_output@invalid_if.nu.snap

@@ -26,3 +26,8 @@ snapshot_kind: text
 6: error
 ==== TYPE ERRORS ====
 Error (NodeId 0): The condition for if branch is not a boolean
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 5): node If { condition: NodeId(0), then_block: NodeId(2), else_block: Some(NodeId(4)) } not suported yet

src/snapshots/new_nu_parser__test__node_output@invalid_types.nu.snap

@@ -66,3 +66,8 @@ snapshot_kind: text
 ==== TYPE ERRORS ====
 Error (NodeId 7): list must have only one type parameter (to allow selection of types, use oneof<int, string> -- WIP)
 Error (NodeId 17): list must have one type parameter
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 13): node Def { name: NodeId(0), params: NodeId(10), return_ty: None, block: NodeId(12) } not suported yet

src/snapshots/new_nu_parser__test__node_output@let_.nu.snap

@@ -19,3 +19,8 @@ snapshot_kind: text
 2: ()
 3: int
 4: int
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(1), is_mutable: false } not suported yet

src/snapshots/new_nu_parser__test__node_output@let_mismatch.nu.snap

@@ -70,3 +70,8 @@ snapshot_kind: text
 ==== TYPE ERRORS ====
 Error (NodeId 13): initializer does not match declared type
 Error (NodeId 26): initializer does not match declared type
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 4): node Let { variable_name: NodeId(0), ty: Some(NodeId(2)), initializer: NodeId(3), is_mutable: false } not suported yet

src/snapshots/new_nu_parser__test__node_output@list.nu.snap

@@ -50,3 +50,8 @@ snapshot_kind: text
 18: string
 19: list<any>
 20: list<any>
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 3): node List([NodeId(0), NodeId(1), NodeId(2)]) not suported yet

src/snapshots/new_nu_parser__test__node_output@literals.nu.snap

@@ -22,3 +22,8 @@ snapshot_kind: text
 4: string
 5: list<any>
 6: list<any>
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 5): node List([NodeId(0), NodeId(1), NodeId(2), NodeId(3), NodeId(4)]) not suported yet

src/snapshots/new_nu_parser__test__node_output@loop.nu.snap

@@ -47,3 +47,8 @@ snapshot_kind: text
 16: unknown
 ==== TYPE ERRORS ====
 Error (NodeId 15): unsupported ast node 'Loop { block: NodeId(14) }' in typechecker
+==== IR ====
+register_count: 0
+file_count: 0
+==== IR ERRORS ====
+Error (NodeId 2): node Let { variable_name: NodeId(0), ty: None, initializer: NodeId(1), is_mutable: true } not suported yet

src/snapshots/new_nu_parser__test__node_output@math.nu.snap

@@ -18,3 +18,10 @@ snapshot_kind: text
 2: int
 3: int
 4: int
+==== IR ====
+register_count: 2
+file_count: 0
+0: LoadLiteral { dst: RegId(0), lit: Int(3) }
+1: LoadLiteral { dst: RegId(1), lit: Int(4) }
+2: BinaryOp { lhs_dst: RegId(0), op: Math(Plus), rhs: RegId(1) }
+3: Return { src: RegId(0) }