Add inlining #ifNil:* and #ifNotNil:* in the BC interpreter

smarr · smarr · commit 73bd707b9a21 · 2025-02-20T21:18:01.000Z
Signed-off-by: Stefan Marr &lt;git@stefan-marr.de&gt;
diff --git a/src/som/compiler/bc/bytecode_generator.py b/src/som/compiler/bc/bytecode_generator.py
@@ -3,6 +3,13 @@
 from som.vm.symbols import sym_nil, sym_true, sym_false
 
 
+class JumpCondition(object):
+    on_nil = 0
+    on_not_nil = 1
+    on_true = 2
+    on_false = 3
+
+
 def emit_inc(mgenc):
     emit1(mgenc, BC.inc, 0)
 
@@ -186,18 +193,27 @@ def emit_push_constant_index(mgenc, lit_index):
     emit2(mgenc, BC.push_constant, lit_index, 1)
 
 
-def emit_jump_on_bool_with_dummy_offset(mgenc, is_if_true, needs_pop):
+def emit_jump_on_with_dummy_offset(mgenc, condition, needs_pop):
     # Remember: true and false seem flipped here.
     # This is because if the test passes, the block is inlined directly.
     # But if the test fails, we need to jump.
     # Thus, an  `#ifTrue:` needs to generated a jump_on_false.
     if needs_pop:
-        bc = BC.jump_on_false_pop if is_if_true else BC.jump_on_true_pop
         stack_effect = -1
     else:
-        bc = BC.jump_on_false_top_nil if is_if_true else BC.jump_on_true_top_nil
         stack_effect = 0
 
+    if condition == JumpCondition.on_true:
+        bc = BC.jump_on_true_pop if needs_pop else BC.jump_on_true_top_nil
+    elif condition == JumpCondition.on_false:
+        bc = BC.jump_on_false_pop if needs_pop else BC.jump_on_false_top_nil
+    elif condition == JumpCondition.on_not_nil:
+        bc = BC.jump_on_not_nil_pop if needs_pop else BC.jump_on_not_nil_top_top
+    elif condition == JumpCondition.on_nil:
+        bc = BC.jump_on_nil_pop if needs_pop else BC.jump_on_nil_top_top
+    else:
+        raise NotImplementedError("Unknown condition: " + str(condition))
+
     emit1(mgenc, bc, stack_effect)
     idx = mgenc.add_bytecode_argument_and_get_index(0)
     mgenc.add_bytecode_argument(0)
diff --git a/src/som/compiler/bc/method_generation_context.py b/src/som/compiler/bc/method_generation_context.py
@@ -1,12 +1,13 @@
 from rlib.debug import make_sure_not_resized
 from som.compiler.bc.bytecode_generator import (
-    emit_jump_on_bool_with_dummy_offset,
+    emit_jump_on_with_dummy_offset,
     emit_jump_with_dummy_offset,
     emit_pop,
     emit_push_constant,
     emit_jump_backward_with_offset,
     emit_inc_field_push,
     emit_return_field,
+    JumpCondition,
 )
 
 from som.compiler.method_generation_context import MethodGenerationContextBase
@@ -583,7 +584,7 @@ def _assemble_field_setter(self, return_candidate):
         arg_idx = self._bytecode[-(pop_len + return_len + 2)]
         return FieldWrite(self.signature, field_idx, arg_idx)
 
-    def inline_if_true_or_if_false(self, parser, is_if_true):
+    def inline_then_branch(self, parser, condition):
         # HACK: We do assume that the receiver on the stack is a boolean,
         # HACK: similar to the IfTrueIfFalseNode.
         # HACK: We don't support anything but booleans at the moment.
@@ -596,8 +597,8 @@ def inline_if_true_or_if_false(self, parser, is_if_true):
 
         self._remove_last_bytecodes(1)  # remove push_block*
 
-        jump_offset_idx_to_skip_true_branch = emit_jump_on_bool_with_dummy_offset(
-            self, is_if_true, False
+        jump_offset_idx_to_skip_true_branch = emit_jump_on_with_dummy_offset(
+            self, condition, False
         )
 
         # TODO: remove the block from the literal list
@@ -624,7 +625,7 @@ def _has_two_literal_block_arguments(self):
             self._last_bytecode_is_one_of(1, PUSH_BLOCK_BYTECODES) != Bytecodes.invalid
         )
 
-    def inline_if_true_false(self, parser, is_if_true):
+    def inline_then_else_branches(self, parser, condition):
         # HACK: We do assume that the receiver on the stack is a boolean,
         # HACK: similar to the IfTrueIfFalseNode.
         # HACK: We don't support anything but booleans at the moment.
@@ -642,8 +643,8 @@ def inline_if_true_false(self, parser, is_if_true):
             to_be_inlined_2,
         ) = self._extract_block_methods_and_remove_bytecodes()
 
-        jump_offset_idx_to_skip_true_branch = emit_jump_on_bool_with_dummy_offset(
-            self, is_if_true, True
+        jump_offset_idx_to_skip_true_branch = emit_jump_on_with_dummy_offset(
+            self, condition, True
         )
 
         self._is_currently_inlining_a_block = True
@@ -698,8 +699,10 @@ def inline_while(self, parser, is_while_true):
         self._is_currently_inlining_a_block = True
         cond_method.inline(self)
 
-        jump_offset_idx_to_skip_loop_body = emit_jump_on_bool_with_dummy_offset(
-            self, is_while_true, True
+        jump_offset_idx_to_skip_loop_body = emit_jump_on_with_dummy_offset(
+            self,
+            JumpCondition.on_false if is_while_true else JumpCondition.on_true,
+            True,
         )
 
         body_method.inline(self)
@@ -725,8 +728,8 @@ def inline_andor(self, parser, is_or):
 
         self._remove_last_bytecodes(1)  # remove push_block*
 
-        jump_offset_idx_to_skip_branch = emit_jump_on_bool_with_dummy_offset(
-            self, not is_or, True
+        jump_offset_idx_to_skip_branch = emit_jump_on_with_dummy_offset(
+            self, JumpCondition.on_true if is_or else JumpCondition.on_false, True
         )
 
         to_be_inlined = self._literals[block_literal_idx]
diff --git a/src/som/compiler/bc/parser.py b/src/som/compiler/bc/parser.py
@@ -19,7 +19,10 @@
     emit_push_constant_index,
     emit_return_non_local,
 )
-from som.compiler.bc.method_generation_context import MethodGenerationContext
+from som.compiler.bc.method_generation_context import (
+    MethodGenerationContext,
+    JumpCondition,
+)
 from som.compiler.parser import ParserBase
 from som.compiler.symbol import Symbol
 from som.vm.symbols import (
@@ -269,10 +272,21 @@ def _keyword_message(self, mgenc):
 
         if not is_super_send:
             if num_args == 1 and (
-                (keyword == "ifTrue:" and mgenc.inline_if_true_or_if_false(self, True))
+                (
+                    keyword == "ifTrue:"
+                    and mgenc.inline_then_branch(self, JumpCondition.on_false)
+                )
                 or (
                     keyword == "ifFalse:"
-                    and mgenc.inline_if_true_or_if_false(self, False)
+                    and mgenc.inline_then_branch(self, JumpCondition.on_true)
+                )
+                or (
+                    keyword == "ifNil:"
+                    and mgenc.inline_then_branch(self, JumpCondition.on_not_nil)
+                )
+                or (
+                    keyword == "ifNotNil:"
+                    and mgenc.inline_then_branch(self, JumpCondition.on_nil)
                 )
                 or (keyword == "whileTrue:" and mgenc.inline_while(self, True))
                 or (keyword == "whileFalse:" and mgenc.inline_while(self, False))
@@ -284,11 +298,19 @@ def _keyword_message(self, mgenc):
             if num_args == 2 and (
                 (
                     keyword == "ifTrue:ifFalse:"
-                    and mgenc.inline_if_true_false(self, True)
+                    and mgenc.inline_then_else_branches(self, JumpCondition.on_false)
                 )
                 or (
                     keyword == "ifFalse:ifTrue:"
-                    and mgenc.inline_if_true_false(self, False)
+                    and mgenc.inline_then_else_branches(self, JumpCondition.on_true)
+                )
+                or (
+                    keyword == "ifNil:ifNotNil:"
+                    and mgenc.inline_then_else_branches(self, JumpCondition.on_not_nil)
+                )
+                or (
+                    keyword == "ifNotNil:ifNil:"
+                    and mgenc.inline_then_else_branches(self, JumpCondition.on_nil)
                 )
             ):
                 return
diff --git a/src/som/interpreter/bc/bytecodes.py b/src/som/interpreter/bc/bytecodes.py
@@ -80,13 +80,21 @@ class Bytecodes(object):
     jump_on_false_top_nil = jump_on_true_top_nil + 1
     jump_on_true_pop = jump_on_false_top_nil + 1
     jump_on_false_pop = jump_on_true_pop + 1
-    jump_backward = jump_on_false_pop + 1
+    jump_on_not_nil_top_top = jump_on_false_pop + 1
+    jump_on_nil_top_top = jump_on_not_nil_top_top + 1
+    jump_on_not_nil_pop = jump_on_nil_top_top + 1
+    jump_on_nil_pop = jump_on_not_nil_pop + 1
+    jump_backward = jump_on_nil_pop + 1
     jump2 = jump_backward + 1
     jump2_on_true_top_nil = jump2 + 1
     jump2_on_false_top_nil = jump2_on_true_top_nil + 1
     jump2_on_true_pop = jump2_on_false_top_nil + 1
     jump2_on_false_pop = jump2_on_true_pop + 1
-    jump2_backward = jump2_on_false_pop + 1
+    jump2_on_not_nil_top_top = jump2_on_false_pop + 1
+    jump2_on_nil_top_top = jump2_on_not_nil_top_top + 1
+    jump2_on_not_nil_pop = jump2_on_nil_top_top + 1
+    jump2_on_nil_pop = jump2_on_not_nil_pop + 1
+    jump2_backward = jump2_on_nil_pop + 1
 
     q_super_send_1 = jump2_backward + 1
     q_super_send_2 = q_super_send_1 + 1
@@ -154,12 +162,20 @@ def is_one_of(bytecode, candidates):
     Bytecodes.jump_on_true_pop,
     Bytecodes.jump_on_false_pop,
     Bytecodes.jump_on_false_top_nil,
+    Bytecodes.jump_on_not_nil_top_top,
+    Bytecodes.jump_on_nil_top_top,
+    Bytecodes.jump_on_not_nil_pop,
+    Bytecodes.jump_on_nil_pop,
     Bytecodes.jump_backward,
     Bytecodes.jump2,
     Bytecodes.jump2_on_true_top_nil,
     Bytecodes.jump2_on_true_pop,
     Bytecodes.jump2_on_false_pop,
     Bytecodes.jump2_on_false_top_nil,
+    Bytecodes.jump2_on_not_nil_top_top,
+    Bytecodes.jump2_on_nil_top_top,
+    Bytecodes.jump2_on_not_nil_pop,
+    Bytecodes.jump2_on_nil_pop,
     Bytecodes.jump2_backward,
 ]
 
@@ -258,12 +274,20 @@ def is_one_of(bytecode, candidates):
     LEN_TWO_ARGS,  # jump_on_false_top_nil
     LEN_TWO_ARGS,  # jump_on_true_pop
     LEN_TWO_ARGS,  # jump_on_false_pop
+    LEN_TWO_ARGS,  # jump_on_not_nil_top_top,
+    LEN_TWO_ARGS,  # jump_on_nil_top_top,
+    LEN_TWO_ARGS,  # jump_on_not_nil_pop,
+    LEN_TWO_ARGS,  # jump_on_nil_pop,
     LEN_TWO_ARGS,  # jump_backward
     LEN_TWO_ARGS,  # jump2
     LEN_TWO_ARGS,  # jump2_on_true_top_nil
     LEN_TWO_ARGS,  # jump2_on_false_top_nil
     LEN_TWO_ARGS,  # jump2_on_true_pop
     LEN_TWO_ARGS,  # jump2_on_false_pop
+    LEN_TWO_ARGS,  # jump2_on_not_nil_top_top,
+    LEN_TWO_ARGS,  # jump2_on_nil_top_top,
+    LEN_TWO_ARGS,  # jump2_on_not_nil_pop,
+    LEN_TWO_ARGS,  # jump2_on_nil_pop,
     LEN_TWO_ARGS,  # jump2_backward
     LEN_ONE_ARG,  # q_super_send_1
     LEN_ONE_ARG,  # q_super_send_2
diff --git a/src/som/interpreter/bc/interpreter.py b/src/som/interpreter/bc/interpreter.py
@@ -584,6 +584,48 @@ def interpret(method, frame, max_stack_size):
                 stack[stack_ptr] = None
             stack_ptr -= 1
 
+        elif bytecode == Bytecodes.jump_on_not_nil_top_top:
+            val = stack[stack_ptr]
+            if val is not nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1)
+                # stack[stack_ptr] = val
+            else:
+                if we_are_jitted():
+                    stack[stack_ptr] = None
+                stack_ptr -= 1
+                current_bc_idx += LEN_TWO_ARGS
+
+        elif bytecode == Bytecodes.jump_on_nil_top_top:
+            val = stack[stack_ptr]
+            if val is nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1)
+                # stack[stack_ptr] = val
+            else:
+                if we_are_jitted():
+                    stack[stack_ptr] = None
+                stack_ptr -= 1
+                current_bc_idx += LEN_TWO_ARGS
+
+        elif bytecode == Bytecodes.jump_on_not_nil_pop:
+            val = stack[stack_ptr]
+            if val is not nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1)
+            else:
+                current_bc_idx += LEN_TWO_ARGS
+            if we_are_jitted():
+                stack[stack_ptr] = None
+            stack_ptr -= 1
+
+        elif bytecode == Bytecodes.jump_on_nil_pop:
+            val = stack[stack_ptr]
+            if val is nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1)
+            else:
+                current_bc_idx += LEN_TWO_ARGS
+            if we_are_jitted():
+                stack[stack_ptr] = None
+            stack_ptr -= 1
+
         elif bytecode == Bytecodes.jump_backward:
             current_bc_idx -= method.get_bytecode(current_bc_idx + 1)
             jitdriver.can_enter_jit(
@@ -649,6 +691,55 @@ def interpret(method, frame, max_stack_size):
                 stack[stack_ptr] = None
             stack_ptr -= 1
 
+        elif bytecode == Bytecodes.jump2_on_not_nil_top_top:
+            val = stack[stack_ptr]
+            if val is not nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1) + (
+                    method.get_bytecode(current_bc_idx + 2) << 8
+                )
+                # stack[stack_ptr] = val
+            else:
+                if we_are_jitted():
+                    stack[stack_ptr] = None
+                stack_ptr -= 1
+                current_bc_idx += LEN_TWO_ARGS
+
+        elif bytecode == Bytecodes.jump2_on_nil_top_top:
+            val = stack[stack_ptr]
+            if val is nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1) + (
+                    method.get_bytecode(current_bc_idx + 2) << 8
+                )
+                # stack[stack_ptr] = val
+            else:
+                if we_are_jitted():
+                    stack[stack_ptr] = None
+                stack_ptr -= 1
+                current_bc_idx += LEN_TWO_ARGS
+
+        elif bytecode == Bytecodes.jump2_on_not_nil_pop:
+            val = stack[stack_ptr]
+            if val is not nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1) + (
+                    method.get_bytecode(current_bc_idx + 2) << 8
+                )
+            else:
+                current_bc_idx += LEN_TWO_ARGS
+            if we_are_jitted():
+                stack[stack_ptr] = None
+            stack_ptr -= 1
+
+        elif bytecode == Bytecodes.jump2_on_nil_pop:
+            val = stack[stack_ptr]
+            if val is nilObject:
+                current_bc_idx += method.get_bytecode(current_bc_idx + 1) + (
+                    method.get_bytecode(current_bc_idx + 2) << 8
+                )
+            else:
+                current_bc_idx += LEN_TWO_ARGS
+            if we_are_jitted():
+                stack[stack_ptr] = None
+            stack_ptr -= 1
         elif bytecode == Bytecodes.jump2_backward:
             current_bc_idx -= method.get_bytecode(current_bc_idx + 1) + (
                 method.get_bytecode(current_bc_idx + 2) << 8
diff --git a/src/som/vmobjects/method_bc.py b/src/som/vmobjects/method_bc.py
diff --git a/tests/test_bytecode_generation.py b/tests/test_bytecode_generation.py
