Merge pull request #16 from Rigidity/inline-definition-fix

Fix inline function inner definitions
xch-dev · Jun 26, 2024 · e519668 · e519668
2 parents 22514ef + d5a3e7a
commit e519668
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 9 deletions.
diff --git a/crates/rue-compiler/src/optimizer.rs b/crates/rue-compiler/src/optimizer.rs
@@ -52,6 +52,12 @@ impl<'a> Optimizer<'a> {
     }
 
     fn opt_path(&mut self, env_id: EnvironmentId, symbol_id: SymbolId) -> LirId {
+        for inline_parameter_map in self.inline_parameter_stack.iter().rev() {
+            if let Some(lir_id) = inline_parameter_map.get(&symbol_id) {
+                return *lir_id;
+            }
+        }
+
         let mut current_env_id = env_id;
         let mut environment = self.db.env(env_id).build().clone();
 
@@ -66,7 +72,12 @@ impl<'a> Optimizer<'a> {
         let index = environment
             .iter()
             .position(|&id| id == symbol_id)
-            .unwrap_or_else(|| panic!("symbol not found in environment: {symbol_id:?}"));
+            .unwrap_or_else(|| {
+                panic!(
+                    "Symbol `{}` not found in environment.",
+                    self.db.dbg_symbol(symbol_id)
+                );
+            });
 
         let mut path = 1;
 
@@ -308,13 +319,7 @@ impl<'a> Optimizer<'a> {
             Symbol::Let(symbol) if self.graph.symbol_usages(symbol_id) == 1 => {
                 self.opt_hir(env_id, symbol.hir_id)
             }
-            Symbol::Let(..) | Symbol::Const(..) => self.opt_path(env_id, symbol_id),
-            Symbol::Parameter { .. } => {
-                for inline_parameter_map in self.inline_parameter_stack.iter().rev() {
-                    if let Some(lir_id) = inline_parameter_map.get(&symbol_id) {
-                        return *lir_id;
-                    }
-                }
+            Symbol::Let(..) | Symbol::Const(..) | Symbol::Parameter(..) => {
                 self.opt_path(env_id, symbol_id)
             }
             Symbol::Unknown | Symbol::Module(..) => unreachable!(),
@@ -371,6 +376,20 @@ impl<'a> Optimizer<'a> {
         let mut inline_parameter_map = HashMap::new();
         let mut args = args;
 
+        let env = self.db.env(function_env_id).clone();
+        for symbol_id in env.definitions() {
+            if self.db.symbol(symbol_id).is_parameter() {
+                continue;
+            }
+            self.db.env_mut(env_id).define(symbol_id);
+        }
+        for symbol_id in env.captures() {
+            if self.db.symbol(symbol_id).is_parameter() {
+                continue;
+            }
+            self.db.env_mut(env_id).capture(symbol_id);
+        }
+
         let param_len = self.db.env(function_env_id).parameters().len();
 
         for (i, symbol_id) in self

diff --git a/crates/rue-compiler/src/optimizer/environment.rs b/crates/rue-compiler/src/optimizer/environment.rs
@@ -2,7 +2,7 @@ use indexmap::IndexSet;
 
 use crate::{EnvironmentId, SymbolId};
 
-#[derive(Debug, Default)]
+#[derive(Debug, Clone, Default)]
 pub struct Environment {
     defined_symbols: IndexSet<SymbolId>,
     captured_symbols: IndexSet<SymbolId>,

diff --git a/tests.toml b/tests.toml
@@ -449,3 +449,10 @@ cost = 87452
 input = "((0x4696e7a2b7682e2df01ab47e6e002d0dca895f99c6172e4a55a3e033499532b7 0x32ed6e4964102d7093ef350019dfd14c99a3ea9feac1f3502194384b8976f7c1 0x770bb2c0b2582924adf403e62374f8424a2ed510eef70b5f450eccab238a4911) 0x8d496dc7cdbc417db2132eda6894b29a91511f176e93a1ea943d210cd27822b2 0x3878fc2ed6b703c7c3d00f9f5d8f170ec252ca439be6e5b5b516e1dce1adc0d7 0x3fe4d62a7db919b25377cb207f16fa0fb6519e123366eaa23014dd5d7c967ca2 0x837525fb91b88dbf02d9ef5a0b322f6943f93424b6e8fe6d26dd1be6d3311871 1 0xc19110971a0cea01368f0c7599b8984f74e301b7cda20bd7f86eb2296bbf16c5)"
 output = "((73 1) (72 0x6155f55414a5cd9193bef33744095a78d57852cf24241b25edeffad6e544c499) (63 0x9d6824bfdfb4c726334a210b657a5e4126c3fbb378598bb3a5b7a210bb75cdb8) (g1_negate 0xc19110971a0cea01368f0c7599b8984f74e301b7cda20bd7f86eb2296bbf16c5 1 (0xc19110971a0cea01368f0c7599b8984f74e301b7cda20bd7f86eb2296bbf16c5)) (60 0x81fdd3fbc407906bc0875522e7a2e77409ee5ef17d3eaa611b505b344588f6b6))"
 hash = "4696e7a2b7682e2df01ab47e6e002d0dca895f99c6172e4a55a3e033499532b7"
+
+[inner_inline_functions]
+bytes = 338
+cost = 16974
+input = "()"
+output = "0x0a37e3fbe2"
+hash = "09c547472ecee6e0a98aee751c1f8ce575c973578acf81c096e4b92e98ddad8b"
diff --git a/tests/functions/inner_inline_functions.rue b/tests/functions/inner_inline_functions.rue
@@ -0,0 +1,45 @@
+const CONST: Int = 43874913874;
+
+fun main() -> Int {
+    outer_1(42) + outer_2(42) + outer_3(42) + outer_4(42) + outer_closure(42)(42)
+}
+
+inline fun outer_closure(num_1: Int) -> fun(num: Int) -> Int {
+    const FACTOR: Int = 2;
+    fun inner(num_2: Int) -> Int {
+        num_1 * num_2 * FACTOR * outer_1(num_1 * num_2) + CONST
+    }
+    inner
+}
+
+inline fun outer_1(num: Int) -> Int {
+    inline const FACTOR: Int = 2;
+    inline fun inner() -> Int {
+        num * FACTOR
+    }
+    inner()
+}
+
+inline fun outer_2(num: Int) -> Int {
+    const FACTOR: Int = 2;
+    inline fun inner() -> Int {
+        num * FACTOR
+    }
+    inner()
+}
+
+inline fun outer_3(num: Int) -> Int {
+    const FACTOR: Int = 2;
+    fun inner() -> Int {
+        num * FACTOR
+    }
+    inner()
+}
+
+inline fun outer_4(num: Int) -> Int {
+    inline const FACTOR: Int = 2;
+    fun inner() -> Int {
+        num * FACTOR
+    }
+    inner()
+}