bridge: inject static box singleton for LLVM instruction ModuleFunction calls

src/runner/modes/common_util/core_bridge.rs

@@ -24,8 +24,8 @@ pub fn canonicalize_module_json(input: &str) -> Result<String, String> {
         }
     }
 
-    let inject_singleton = env_flag("HAKO_BRIDGE_INJECT_SINGLETON")
+    let inject_singleton =
-        || env_flag("NYASH_BRIDGE_INJECT_SINGLETON");
+        env_flag("HAKO_BRIDGE_INJECT_SINGLETON") || env_flag("NYASH_BRIDGE_INJECT_SINGLETON");
     let materialize_phi = env_flag("HAKO_BRIDGE_EARLY_PHI_MATERIALIZE")
         || env_flag("NYASH_BRIDGE_EARLY_PHI_MATERIALIZE");
 
@@ -68,14 +68,10 @@ fn inject_singleton_methods(root: &mut Value) -> Result<bool, String> {
     };
 
     for func in functions.iter_mut() {
-        let blocks = func
+        let blocks = func.get_mut("blocks").and_then(Value::as_array_mut);
-            .get_mut("blocks")
-            .and_then(Value::as_array_mut);
         let Some(blocks) = blocks else { continue };
         for block in blocks.iter_mut() {
-            let insts = block
+            let insts = block.get_mut("instructions").and_then(Value::as_array_mut);
-                .get_mut("instructions")
-                .and_then(Value::as_array_mut);
             let Some(insts) = insts else { continue };
             for inst in insts.iter_mut() {
                 if transform_module_function(inst)? {
@@ -118,12 +114,26 @@ fn transform_module_function(inst: &mut Value) -> Result<bool, String> {
     };
     let name = name_owned.as_str();
 
-    let (box_name, method) = match name {
+    match name {
-        "ArrayBox.len" => ("ArrayBox", "size"),
+        "ArrayBox.len" | "MapBox.len" => {
-        "MapBox.len" => ("MapBox", "len"),
+            rewrite_array_like_length(mir_call, name)?;
-        _ => return Ok(false),
+            Ok(true)
-    };
+        }
+        _ => {
+            if let Some((box_name, method)) = static_singleton_method(name) {
+                rewrite_static_singleton_call(mir_call, box_name, method)?;
+                Ok(true)
+            } else {
+                Ok(false)
+            }
+        }
+    }
+}
 
+fn rewrite_array_like_length(
+    mir_call: &mut serde_json::Map<String, Value>,
+    name: &str,
+) -> Result<(), String> {
     let receiver_value = mir_call
         .get_mut("args")
         .and_then(Value::as_array_mut)
@@ -146,6 +156,11 @@ fn transform_module_function(inst: &mut Value) -> Result<bool, String> {
         .get_mut("callee")
         .and_then(Value::as_object_mut)
         .ok_or_else(|| "bridge canonicalize: callee missing".to_string())?;
+    let (box_name, method) = match name {
+        "ArrayBox.len" => ("ArrayBox", "size"),
+        "MapBox.len" => ("MapBox", "len"),
+        _ => unreachable!(),
+    };
     callee.insert("type".to_string(), Value::String("Method".into()));
     callee.insert("method".to_string(), Value::String(method.into()));
     callee.insert("box_name".to_string(), Value::String(box_name.into()));
@@ -155,7 +170,103 @@ fn transform_module_function(inst: &mut Value) -> Result<bool, String> {
         callee.insert("certainty".to_string(), Value::String("Known".into()));
     }
 
-    Ok(true)
+    Ok(())
+}
+
+const STATIC_SINGLETON_METHODS: &[(&str, &str)] = &[
+    ("PhiInst", "lower_phi"),
+    ("ConstInst", "lower_const"),
+    ("BinOpInst", "lower_binop"),
+    ("CompareInst", "lower_compare"),
+    ("BranchInst", "lower_branch"),
+    ("JumpInst", "lower_jump"),
+    ("ReturnInst", "lower_return"),
+    ("LLVMPhiInstructionBox", "lower_phi"),
+    ("LLVMConstInstructionBox", "lower_const"),
+    ("LLVMBinOpInstructionBox", "lower_binop"),
+    ("LLVMCompareInstructionBox", "lower_compare"),
+    ("LLVMBranchInstructionBox", "lower_branch"),
+    ("LLVMJumpInstructionBox", "lower_jump"),
+    ("LLVMReturnInstructionBox", "lower_return"),
+];
+
+fn static_singleton_method(name: &str) -> Option<(&'static str, &'static str)> {
+    STATIC_SINGLETON_METHODS
+        .iter()
+        .find(|(box_name, method)| name == format!("{}.{}", box_name, method))
+        .copied()
+}
+
+fn rewrite_static_singleton_call(
+    mir_call: &mut serde_json::Map<String, Value>,
+    box_name: &'static str,
+    method: &'static str,
+) -> Result<(), String> {
+    let callee = mir_call
+        .get_mut("callee")
+        .and_then(Value::as_object_mut)
+        .ok_or_else(|| "bridge canonicalize: callee missing".to_string())?;
+
+    callee.insert("type".to_string(), Value::String("Method".into()));
+    callee.insert("box_name".to_string(), Value::String(box_name.into()));
+    callee.insert("method".to_string(), Value::String(method.into()));
+    callee.remove("name");
+    callee.remove("module");
+    callee.remove("receiver");
+    if !callee.contains_key("certainty") {
+        callee.insert("certainty".to_string(), Value::String("Known".into()));
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::env;
+
+    fn with_env<F: FnOnce() -> R, R>(key: &str, val: &str, f: F) -> R {
+        let prev = env::var(key).ok();
+        env::set_var(key, val);
+        let result = f();
+        if let Some(prev_val) = prev {
+            env::set_var(key, prev_val);
+        } else {
+            env::remove_var(key);
+        }
+        result
+    }
+
+    #[test]
+    fn singleton_injection_disabled_noop() {
+        env::remove_var("HAKO_BRIDGE_INJECT_SINGLETON");
+        env::remove_var("NYASH_BRIDGE_INJECT_SINGLETON");
+        let input = r#"{"functions":[{"blocks":[{"instructions":[{"op":"mir_call","mir_call":{"callee":{"type":"ModuleFunction","name":"LLVMPhiInstructionBox.lower_phi"},"args":[1,2]}}]}]}]}"#;
+        let output = canonicalize_module_json(input).expect("canonicalize");
+        assert_eq!(output, input);
+    }
+
+    #[test]
+    fn singleton_injection_enabled_rewrites_static_box() {
+        let input = r#"{"functions":[{"blocks":[{"instructions":[{"op":"mir_call","mir_call":{"callee":{"type":"ModuleFunction","name":"LLVMPhiInstructionBox.lower_phi"},"args":[1,2]}}]}]}]}"#;
+        let output = with_env("HAKO_BRIDGE_INJECT_SINGLETON", "1", || {
+            canonicalize_module_json(input).expect("canonicalize")
+        });
+        let value: Value = serde_json::from_str(&output).expect("json");
+        let callee = value["functions"][0]["blocks"][0]["instructions"][0]["mir_call"]["callee"]
+            .as_object()
+            .expect("callee object");
+        assert_eq!(callee.get("type").and_then(Value::as_str), Some("Method"));
+        assert_eq!(
+            callee.get("box_name").and_then(Value::as_str),
+            Some("LLVMPhiInstructionBox")
+        );
+        assert_eq!(
+            callee.get("method").and_then(Value::as_str),
+            Some("lower_phi")
+        );
+        assert!(callee.get("name").is_none());
+    }
 }
 
 fn materialize_phi_blocks(root: &mut Value) -> Result<bool, String> {
@@ -171,14 +282,10 @@ fn materialize_phi_blocks(root: &mut Value) -> Result<bool, String> {
     };
 
     for func in functions.iter_mut() {
-        let blocks = func
+        let blocks = func.get_mut("blocks").and_then(Value::as_array_mut);
-            .get_mut("blocks")
-            .and_then(Value::as_array_mut);
         let Some(blocks) = blocks else { continue };
         for block in blocks.iter_mut() {
-            let insts = block
+            let insts = block.get_mut("instructions").and_then(Value::as_array_mut);
-                .get_mut("instructions")
-                .and_then(Value::as_array_mut);
             let Some(insts) = insts else { continue };
             if reorder_block_phi(insts)? {
                 changed = true;