Improve infer

compiler/src/infer.rs

@@ -12,7 +12,6 @@ use crate::{
 pub enum FnKind {
     Id(usize),
     DefineField(AtomRef),
-    PutField(AtomRef),
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -42,16 +41,44 @@ fn join_ty(a: &Type, b: &Type) -> Type {
     }
 }
 
+fn get_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Type {
+    match ty {
+        Type::ReturnOf((fn_kind, args)) => match fn_kind {
+            FnKind::Id(id) => {
+                let ret_ty = cache
+                    .get(&InferRequest {
+                        fn_id: *id,
+                        arg_types: args.clone(),
+                    })
+                    .unwrap()
+                    .ret_type
+                    .clone();
+
+                get_field(&ret_ty, field, cache)
+            }
+            FnKind::DefineField(atom) => {
+                if *atom == field {
+                    args[0].clone()
+                } else {
+                    Type::Undefined
+                }
+            }
+        },
+        Type::Any => Type::Any,
+        _ => Type::Undefined,
+    }
+}
+
 #[derive(Debug, Hash, PartialEq, Eq)]
 pub struct InferRequest {
-    fn_id: usize,
+    pub fn_id: usize,
-    arg_types: Vec<Type>,
+    pub arg_types: Vec<Type>,
 }
 
 #[derive(Debug)]
 pub struct InferResult {
-    val_types: Vec<Type>,
+    pub val_types: Vec<Type>,
-    ret_type: Type,
+    pub ret_type: Type,
 }
 
 pub type TypeCache = HashMap<InferRequest, InferResult>;
@@ -64,55 +91,44 @@ pub fn infer_types(reg_fns: &[RegFn], atoms: &[Atom]) -> TypeCache {
         .unwrap();
 
     let mut cache = TypeCache::with_capacity(reg_fns.len());
-    let mut worklist = vec![];
 
-    worklist.push(InferRequest {
+    let req = InferRequest {
         fn_id: main_fn_id,
         arg_types: vec![],
-    });
+    };
 
-    while let Some(req) = worklist.pop() {
+    walk_fn(req, reg_fns, &mut cache);
-        match walk_fn(&req, reg_fns, &mut cache) {
-            Ok(res) => {
-                cache.insert(req, res);
-            }
-            Err(new_req) => {
-                worklist.push(new_req);
-                worklist.push(req);
-            }
-        }
-    }
 
     cache
 }
 
-fn walk_fn(
+fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
-    req: &InferRequest,
-    reg_fns: &[RegFn],
-    cache: &mut TypeCache,
-) -> Result<InferResult, InferRequest> {
     use OpCode::*;
 
     let f = &reg_fns[req.fn_id];
     let mut val_types = vec![Type::Unknown; f.val_count];
     let mut ret_type = Type::Unknown;
 
-    for instr in &f.code {
+    for block in &f.blocks {
+        let code = &f.code[block.start..block.end];
+        let code = code.iter().chain(&block.postfix);
+
+        for instr in code {
             let is_ret = matches!(instr.op, Return | ReturnAsync | ReturnUndef);
 
             if instr.dsts.is_empty() && !is_ret {
                 continue;
             }
 
+            // TODO: Some of these ops may not every show up by this point
             let ty = match instr.op {
                 // ──────────────────────────────────── Literals ────────────────────────────────────
-            PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7 | PushI8
+                PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7
-            | PushI16 | PushI32 => Type::Int,
+                | PushI8 | PushI16 | PushI32 => Type::Int,
 
                 PushTrue | PushFalse => Type::Bool,
 
-            // TODO: Fclosure will disappear soon
+                PushConst | PushConst8 => match instr.imm {
-            PushConst | PushConst8 | Fclosure | Fclosure8 => match instr.imm {
                     ImmOperand::Const(k) => match f.cpool.get(k as usize) {
                         Some(JsValue::Int(_)) => Type::Int,
                         Some(JsValue::Float(_)) => Type::Float,
@@ -124,7 +140,7 @@ fn walk_fn(
                 },
 
                 // ─────────────────────────────────── Arithmetic ───────────────────────────────────
-            Add | Sub | Mul | Div | Mod | Pow => {
+                Add | Sub | Mul | Mod | Pow => {
                     match (&val_types[instr.args[0]], &val_types[instr.args[1]]) {
                         (Type::Int, Type::Int) => Type::Int,
                         (Type::Float, Type::Float)
@@ -133,25 +149,24 @@ fn walk_fn(
                         _ => Type::Any,
                     }
                 }
-
                 Neg | Plus | Inc | Dec => match val_types[instr.args[0]] {
                     Type::Int => Type::Int,
                     Type::Float => Type::Float,
                     _ => Type::Any,
                 },
+                Div => Type::Float,
 
                 Lt | Lte | Gt | Gte | Eq | Neq | StrictEq | StrictNeq | Not => Type::Bool,
 
                 // ────────────────────────────────── Value Movers ──────────────────────────────────
-            Dup | SetLoc => val_types[instr.args[0]].clone(),
+                Dup | SetLoc | SetName => val_types[instr.args[0]].clone(),
 
                 // ──────────────────────────────────── Arguments ───────────────────────────────────
-            OpCode::GetArg0 => req.arg_types.first().cloned().unwrap_or(Type::Undefined),
+                GetArg0 => req.arg_types.first().cloned().unwrap_or(Type::Undefined),
-            OpCode::GetArg1 => req.arg_types.get(1).cloned().unwrap_or(Type::Undefined),
+                GetArg1 => req.arg_types.get(1).cloned().unwrap_or(Type::Undefined),
-            OpCode::GetArg2 => req.arg_types.get(2).cloned().unwrap_or(Type::Undefined),
+                GetArg2 => req.arg_types.get(2).cloned().unwrap_or(Type::Undefined),
-            OpCode::GetArg3 => req.arg_types.get(3).cloned().unwrap_or(Type::Undefined),
+                GetArg3 => req.arg_types.get(3).cloned().unwrap_or(Type::Undefined),
-
+                GetArg => match instr.imm {
-            OpCode::GetArg => match instr.imm {
                     ImmOperand::Arg(k) => req
                         .arg_types
                         .get(k as usize)
@@ -161,9 +176,9 @@ fn walk_fn(
                 },
 
                 // ──────────────────────────────────── Objects ─────────────────────────────────────
-            OpCode::Object => Type::Empty,
+                Object => Type::Empty,
-            OpCode::DefineField => {
+                DefineField => {
-                let atom = match instr.imm {
+                    let field = match instr.imm {
                         ImmOperand::Atom(atom) => atom,
                         _ => unreachable!(),
                     };
@@ -171,14 +186,26 @@ fn walk_fn(
                     let arg0_ty = val_types[instr.args[0]].clone();
                     let arg1_ty = val_types[instr.args[1]].clone();
 
-                ReturnOf((FnKind::DefineField(atom), vec![arg0_ty, arg1_ty]))
+                    ReturnOf((FnKind::DefineField(field), vec![arg0_ty, arg1_ty]))
                 }
+                GetField => {
+                    let field = match instr.imm {
+                        ImmOperand::Atom(atom) => atom,
+                        _ => unreachable!(),
+                    };
+
+                    let ty = val_types[instr.args[0]].clone();
+                    get_field(&ty, field, cache)
+                }
+
+                // ──────────────────────────────── Function Calls ──────────────────────────────────
+                Call | Call0 | Call1 | Call2 | Call3 | TailCall => todo!(),
 
                 // ────────────────────────────────────── Returns ───────────────────────────────────
                 Return | ReturnAsync => val_types[instr.args[0]].clone(),
                 ReturnUndef => Type::Empty,
 
-            _ => unreachable!(),
+                _ => unreachable!("{instr:?}"),
             };
 
             if is_ret {
@@ -189,14 +216,18 @@ fn walk_fn(
                 }
             }
         }
+    }
 
-    Ok(InferResult {
+    cache.insert(
+        req,
+        InferResult {
             val_types,
             ret_type,
-    })
+        },
+    );
 }
 
-pub fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
+fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
     match a {
         AtomRef::TaggedInt(n) => n.to_string(),
         AtomRef::Predefined(i) => atom_names::PREDEFINED_ATOMS