feat(closure): extract top level vars

src/bin/main.mbt

@@ -67,6 +67,7 @@ struct CompileStatus {
   mut asm : @riscv.AssemblyProgram?
   mut js : @js.JsResult?
   mut wasm : @wasm.WasmResult?
+  mut optimizing_manager : @optimizing.Manager?
 }
 
 fn CompileStatus::initialize(
@@ -89,6 +90,7 @@ fn CompileStatus::initialize(
     asm: None,
     js: None,
     wasm: None,
+    optimizing_manager: None,
   }
   match start_stage {
     Parse => v.source_code = Some(file)
@@ -120,7 +122,9 @@ fn CompileStatus::step(self : CompileStatus) -> Bool {
     }
     Typecheck => {
       let typechecked = @typing.type_check(self.ast.unwrap(), externals())
-      let typechecked = optimizing_manager.optimize_ast(typechecked)
+      let typechecked = self.optimizing_manager
+        .unwrap()
+        .optimize_ast(typechecked)
       self.typechecked = Some(typechecked)
     }
     Knf => {
@@ -133,7 +137,7 @@ fn CompileStatus::step(self : CompileStatus) -> Bool {
     }
     KnfOpt => {
       let knf = self.knf.unwrap()
-      let knf = optimizing_manager.optimize_knf(knf)
+      let knf = self.optimizing_manager.unwrap().optimize_knf(knf)
       self.opt_knf = Some(knf)
     }
     Closure => {
@@ -149,17 +153,17 @@ fn CompileStatus::step(self : CompileStatus) -> Bool {
       let world = builder.build(self.closure_ir.unwrap())
       self.ssa = Some(world)
       // Optimizing SSA
-      optimizing_manager.optimize_ssa_ir(world)
+      self.optimizing_manager.unwrap().optimize_ssa_ir(world)
     }
     Asm => {
       let world = self.ssa.unwrap()
       // Backend
       match self.backend {
         RSICV => {
           let machine_world = @riscv.build_machine_world(world)
-          optimizing_manager.optimize_machine_ir(machine_world)
+          self.optimizing_manager.unwrap().optimize_machine_ir(machine_world)
           let asm_program = @riscv.emit(machine_world)
-          optimizing_manager.optimize_assembly(asm_program)
+          self.optimizing_manager.unwrap().optimize_assembly(asm_program)
           self.asm = Some(asm_program)
         }
         JS => self.js = Some(@js.emit_js(world))
@@ -367,8 +371,12 @@ fn main {
     Ok(status) => {
       if js_backend.val {
         status.backend = JS
+        status.optimizing_manager = Some(get_optimizing_manager(JS))
       } else if wasm_backend.val {
         status.backend = WASM
+        status.optimizing_manager = Some(get_optimizing_manager(WASM))
+      } else {
+        status.optimizing_manager = Some(get_optimizing_manager(RSICV))
       }
       status
     }

src/bin/optimizing.mbt

@@ -6,15 +6,22 @@ typealias Pass = @optimizing.Pass
 
 // typealias ClosureOptPass = @optimizing.ClosureOptPass
 
-let optimizing_manager : @optimizing.Manager = @optimizing.new().add_opt_passes(
-  [
-    // Pass::Ast(ConstFold),
-    // Pass::Ast(Liveness),
-    Pass::Knf(CopyPropagation),
-    Pass::Knf(TailRecursion),
-    Pass::Knf(MayTailRecursion),
-    Pass::Ssa(ConstProp),
-    Pass::Ssa(DeadCodeElimination),
-    Pass::Assembly(InstFold),
-  ],
-)
+fn get_optimizing_manager(backend : Backends) -> @optimizing.Manager {
+  let optimizing_manager : @optimizing.Manager = @optimizing.new().add_opt_passes(
+    [
+      // Pass::Ast(ConstFold),
+      // Pass::Ast(Liveness),
+      Pass::Knf(CopyPropagation),
+      Pass::Knf(TailRecursion),
+      Pass::Knf(MayTailRecursion),
+      Pass::Ssa(ConstProp),
+      Pass::Ssa(DeadCodeElimination),
+      Pass::Assembly(InstFold),
+    ],
+  )
+  if backend == RSICV {
+    optimizing_manager.add_opt_pass(Pass::Knf(MethodInline))
+  } else {
+    optimizing_manager
+  }
+}

src/closure/closure_ir.mbt

@@ -40,6 +40,7 @@ pub enum Expr {
 
 pub struct Closure {
   name : Label
+  top_level_vars : Array[(Name, LowType)]
   actual_free_vars : Array[Name]
 } derive(Show)
 
@@ -54,6 +55,7 @@ pub struct FuncDef {
   /// true if the function is a closure function, and the closure will be available at `s11` on 
   /// function entry.
   mut is_closure : Bool
+  mut top_level_vars : Array[(Name, LowType)]
   mut formal_free_vars : Array[(Name, LowType)]
 } derive(Show)
 

src/closure/json.mbt

@@ -197,7 +197,11 @@ pub fn Closure::from_json(json : Json) -> Closure! {
       for n in actual_free_vars {
         actual_free_vars_res.push(Name::from_json!(n))
       }
-      { name: Label(name), actual_free_vars: actual_free_vars_res }
+      {
+        name: Label(name),
+        top_level_vars: [],
+        actual_free_vars: actual_free_vars_res,
+      }
     }
     _ => fail!("invalid json")
   }
@@ -269,6 +273,7 @@ pub fn FuncDef::from_json(json : Json) -> FuncDef! {
         is_closure: @json.from_json!(b),
         ty: LowType::from_json!(t),
         args: args_res,
+        top_level_vars: [],
         formal_free_vars: formal_free_vars_res,
         body: Expr::from_json!(body),
       }

src/closure/knf_preprocess.mbt

@@ -3,29 +3,35 @@
 struct KnfEnv {
   world : KnfWorld
   scope_path : Array[MutFuncDef]
-  global : @immut/hashset.T[Name]
   local : @hashset.T[Name]
+  global : @immut/hashset.T[Name]
   mut free_vars : @immut/hashset.T[Name]
+  mut top_level : @immut/hashmap.T[Name, LowType]
+  mut top_level_free_vars : @immut/hashmap.T[Name, LowType]
 }
 
 fn KnfEnv::new(externals : Map[String, Type]) -> KnfEnv {
   let mut world_global = @immut/hashmap.new()
   let mut global = @immut/hashset.new()
+  let mut top_level = @immut/hashmap.new()
   externals.each(
     fn {
       name, typ => {
         let binding = Name::name_only(@shared.ffi_prefix + name)
         world_global = world_global.add(binding, to_direct_fn_low_type(typ))
         global = global.add(binding)
+        top_level = top_level.add(binding, to_direct_fn_low_type(typ))
       }
     },
   )
   {
     world: KnfWorld::{ externals: world_global, func_defs: Map::new() },
     scope_path: [],
+    local: @hashset.new(),
     global,
+    top_level,
     free_vars: @immut/hashset.new(),
-    local: @hashset.new(),
+    top_level_free_vars: @immut/hashmap.new(),
   }
 }
 
@@ -35,26 +41,39 @@ fn KnfEnv::fork(self : KnfEnv) -> KnfEnv {
   {
     world: self.world,
     scope_path: self.scope_path,
+    local: @hashset.new(),
     global,
+    top_level: self.top_level,
     free_vars: @immut/hashset.new(),
-    local: @hashset.new(),
+    top_level_free_vars: @immut/hashmap.new(),
   }
 }
 
+fn KnfEnv::add_top_level(self : KnfEnv, name : Name, typ : LowType) -> Unit {
+  self.local.insert(name)
+  self.top_level = self.top_level.add(name, typ)
+}
+
 fn KnfEnv::add_local_var(self : KnfEnv, name : Name) -> Unit {
   self.local.insert(name)
 }
 
-fn KnfEnv::collect_free_var(self : KnfEnv, name : Name) -> Unit {
+fn KnfEnv::collect_free_var(
+  self : KnfEnv,
+  name : Name,
+  ~loc : SourceLoc = _
+) -> Unit {
   match self.local.contains(name) {
     true => ()
     false =>
       match self.global.contains(name) {
         true =>
-          if self.world.externals.find(name).is_empty() {
-            self.free_vars = self.free_vars.add(name)
+          match self.top_level.find(name) {
+            Some(typ) =>
+              self.top_level_free_vars = self.top_level_free_vars.add(name, typ)
+            None => self.free_vars = self.free_vars.add(name)
           }
-        false => @util.die("Unknown identifier name \{name}")
+        false => @util.die("Unknown identifier name \{name}", ~loc)
       }
   }
 }
@@ -85,11 +104,11 @@ fn KnfEnv::exit_scope(self : KnfEnv) -> Unit {
 /// -------------
 fn knf_preprocess(knf : Knf, externals : Map[String, Type]) -> (Knf, KnfWorld) {
   let env = KnfEnv::new(externals)
-  let knf = env.visit_stmt(knf)
+  let knf = env.visit_stmt(knf, top_level=true)
   (knf, env.world)
 }
 
-fn KnfEnv::visit_stmt(self : KnfEnv, knf : Knf) -> Knf {
+fn KnfEnv::visit_stmt(self : KnfEnv, knf : Knf, ~top_level : Bool) -> Knf {
   match knf {
     // Literal
     Unit | Int(_) | Double(_) => knf
@@ -122,22 +141,22 @@ fn KnfEnv::visit_stmt(self : KnfEnv, knf : Knf) -> Knf {
     IfEq(lhs, rhs, true_branch, false_branch, typ) => {
       self.collect_free_var(lhs)
       self.collect_free_var(rhs)
-      let true_branch = self.visit_stmt(true_branch)
-      let false_branch = self.visit_stmt(false_branch)
+      let true_branch = self.visit_stmt(true_branch, top_level=false)
+      let false_branch = self.visit_stmt(false_branch, top_level=false)
       IfEq(lhs, rhs, true_branch, false_branch, typ)
     }
     IfLe(lhs, rhs, true_branch, false_branch, typ) => {
       self.collect_free_var(lhs)
       self.collect_free_var(rhs)
-      let true_branch = self.visit_stmt(true_branch)
-      let false_branch = self.visit_stmt(false_branch)
+      let true_branch = self.visit_stmt(true_branch, top_level=false)
+      let false_branch = self.visit_stmt(false_branch, top_level=false)
       IfLe(lhs, rhs, true_branch, false_branch, typ)
     }
 
     // Loop
     Loop(label, args, body) => {
       self.collect_free_vars(args.map(fn { a => a.0 }).iter())
-      Loop(label, args, self.visit_stmt(body))
+      Loop(label, args, self.visit_stmt(body, top_level=false))
     }
     Continue(label, args) => {
       self.collect_free_vars(args.iter())
@@ -169,18 +188,26 @@ fn KnfEnv::visit_stmt(self : KnfEnv, knf : Knf) -> Knf {
 
     // Let
     Let((name, typ), expr, rest) => {
-      let expr = self.visit_stmt(expr)
-      self.add_local_var(name)
-      Let((name, typ), expr, self.visit_stmt(rest))
+      let expr = self.visit_stmt(expr, top_level=false)
+      if top_level {
+        self.add_top_level(name, to_closure_fn_low_type(typ))
+      } else {
+        self.add_local_var(name)
+      }
+      Let((name, typ), expr, self.visit_stmt(rest, ~top_level))
     }
     LetTuple(els, expr, rest) => {
       self.collect_free_var(expr)
       els.each(fn { el => self.add_local_var(el.0) })
-      LetTuple(els, expr, self.visit_stmt(rest))
+      LetTuple(els, expr, self.visit_stmt(rest, ~top_level))
     }
     LetRec(def, rest) => {
       // Add func def to local env
-      self.add_local_var(def.name)
+      if top_level {
+        self.add_top_level(def.name, to_direct_fn_low_type(def.ty))
+      } else {
+        self.add_local_var(def.name)
+      }
       // Create mutable func def to maintain some extra information
       let mutdef = MutFuncDef::new(self, def)
       // 1. Fork new env for function body
@@ -189,16 +216,24 @@ fn KnfEnv::visit_stmt(self : KnfEnv, knf : Knf) -> Knf {
       // 2. Add params to local env
       def.args.each(fn { (name, _) => new_env.add_local_var(name) })
       // 3. Visit the function body, and update the current def
-      let updated_def = mutdef.update_body(new_env.visit_stmt(def.body))
+      let updated_def = mutdef.update_body(
+        new_env.visit_stmt(def.body, top_level=false),
+      )
       // 4. Record function free variables, and nest free variables
       mutdef.free_vars = @immut/hashset.from_array(
         // Note: we should not pass itself for recursion, for it has not been initialized
         new_env.free_vars.iter().filter(fn { n => n != def.name }).collect(),
       )
+      mutdef.top_level_free_vars = @immut/hashmap.from_iter(
+        new_env.top_level_free_vars.iter(),
+      )
       self.collect_free_vars(new_env.free_vars.iter())
+      self.collect_free_vars(
+        new_env.top_level_free_vars.iter().map(fn { v => v.0 }),
+      )
       // 5. Back to the prev env
       new_env.exit_scope()
-      LetRec(updated_def, self.visit_stmt(rest))
+      LetRec(updated_def, self.visit_stmt(rest, ~top_level))
     }
 
     // External

src/closure/knf_to_closure.mbt

@@ -205,6 +205,7 @@ fn Context::visit_stmt(self : Context, knf : Knf, is_top_level : Bool) -> Expr {
         is_closure,
         ty: func_type,
         args: def.args.mapi(fn { i, a => (a.0, args_type[i]) }),
+        top_level_vars: mutdef.top_level_free_vars.iter().collect(),
         formal_free_vars: mutdef.free_vars
         .iter()
         .map(fn { fv => (fv, self.get_low_type_of(fv)) })
@@ -221,6 +222,7 @@ fn Context::visit_stmt(self : Context, knf : Knf, is_top_level : Bool) -> Expr {
           (name, func_type),
           Closure::{
             name: mutdef.to_label(),
+            top_level_vars: mutdef.top_level_free_vars.iter().collect(),
             actual_free_vars: mutdef.free_vars.iter().collect(),
           },
           rest,
@@ -230,6 +232,7 @@ fn Context::visit_stmt(self : Context, knf : Knf, is_top_level : Bool) -> Expr {
           (name, func_type),
           Closure::{
             name: mutdef.to_label(),
+            top_level_vars: mutdef.top_level_free_vars.iter().collect(),
             actual_free_vars: mutdef.free_vars.iter().collect(),
           },
           rest,

src/closure/world.mbt

@@ -7,6 +7,8 @@ struct MutFuncDef {
   force_closure : Bool
   /// Collected free variables
   mut free_vars : @immut/hashset.T[Name]
+  /// Top level
+  mut top_level_free_vars : @immut/hashmap.T[Name, LowType]
 }
 
 fn MutFuncDef::new(env : KnfEnv, def : @knf.FuncDef) -> MutFuncDef {
@@ -16,6 +18,7 @@ fn MutFuncDef::new(env : KnfEnv, def : @knf.FuncDef) -> MutFuncDef {
     def,
     force_closure: false,
     free_vars: @immut/hashset.new(),
+    top_level_free_vars: @immut/hashmap.new(),
   }
   env.world.func_defs.set(def.id, mutdef)
   mutdef