Merge pull request #135 from TritonVM/raw_code_inline

 fix(RawCode): Only inline inlineable code

Author: Sword-Smith

tasm-lib/src/list/higher_order/inner_function.rs

@@ -77,20 +77,43 @@ impl RawCode {
     }
 
     /// Returns `Some(code)` iff the raw code is a function that can be inlined
+    ///
+    /// Type hints and breakpoints are stripped.
     pub fn inlined_body(&self) -> Option<Vec<LabelledInstruction>> {
-        // If the RawCode contains a recurse, it cannot be inlined
-        if self
+        let is_label = |x: &_| matches!(x, LabelledInstruction::Label(_));
+        let is_instruction = |x: &_| matches!(x, LabelledInstruction::Instruction(_));
+        let is_recursive = |x: &_| {
+            matches!(
+                x,
+                LabelledInstruction::Instruction(AnInstruction::Recurse)
+                    | LabelledInstruction::Instruction(AnInstruction::RecurseOrReturn)
+            )
+        };
+
+        if self.function.iter().any(is_recursive) {
+            // recursion needs to be wrapped in a function
+            return None;
+        }
+
+        let mut labels_and_instructions = self
             .function
             .iter()
-            .any(|x| matches!(x, LabelledInstruction::Instruction(AnInstruction::Recurse)))
-        {
-            None
-        } else if self.function.len() == 2 {
-            Some(triton_asm!())
-        } else {
-            // Remove label and `return`
-            Some(self.function[1..=self.function.len() - 2].to_vec())
-        }
+            .filter(|i| is_label(i) || is_instruction(i));
+
+        let Some(first_thing) = labels_and_instructions.next() else {
+            return Some(triton_asm!());
+        };
+        let LabelledInstruction::Label(_) = first_thing else {
+            panic!("Raw Code must start with a label.")
+        };
+
+        let Some(LabelledInstruction::Instruction(AnInstruction::Return)) =
+            labels_and_instructions.next_back()
+        else {
+            panic!("Raw Code is probably buggy: too short, or doesn't end with `return`.");
+        };
+
+        Some(labels_and_instructions.cloned().collect())
     }
 }
 
@@ -219,3 +242,30 @@ impl InnerFunction {
         };
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn breakpoint_does_not_influence_raw_code_inlining() {
+        let raw_code = RawCode {
+            function: triton_asm! { my_label: return break },
+            input_type: DataType::VoidPointer,
+            output_type: DataType::VoidPointer,
+        };
+        let inlined_code = raw_code.inlined_body().unwrap();
+        assert_eq!(triton_asm!(), inlined_code);
+    }
+
+    #[test]
+    fn type_hints_do_not_influence_raw_code_inlining() {
+        let raw_code = RawCode {
+            function: triton_asm! { my_label: hint a = stack[0] hint b = stack[1] return },
+            input_type: DataType::VoidPointer,
+            output_type: DataType::VoidPointer,
+        };
+        let inlined_code = raw_code.inlined_body().unwrap();
+        assert_eq!(triton_asm!(), inlined_code);
+    }
+}