Add inlining for #ifNil:* variants (#65)

Author: smarr

src/som/compiler/ast/parser.py

@@ -24,6 +24,9 @@
 from som.interpreter.ast.nodes.specialized.literal_if import (
     IfInlinedNode,
     IfElseInlinedNode,
+    IfNilInlinedNode,
+    IfNotNilInlinedNode,
+    IfNilNotNilInlinedNode,
 )
 from som.interpreter.ast.nodes.specialized.literal_while import WhileInlinedNode
 from som.vm.symbols import symbol_for
@@ -249,6 +252,26 @@ def _try_inlining_if(if_true, receiver, arguments, source, mgenc):
         body_expr = method.inline(mgenc)
         return IfInlinedNode(receiver, body_expr, if_true, source)
 
+    @staticmethod
+    def _try_inlining_if_nil(receiver, arguments, source, mgenc):
+        arg = arguments[0]
+        if not isinstance(arg, BlockNode):
+            return None
+
+        method = arg.get_method()
+        body_expr = method.inline(mgenc)
+        return IfNilInlinedNode(receiver, body_expr, source)
+
+    @staticmethod
+    def _try_inlining_if_not_nil(receiver, arguments, source, mgenc):
+        arg = arguments[0]
+        if not isinstance(arg, BlockNode):
+            return None
+
+        method = arg.get_method()
+        body_expr = method.inline(mgenc)
+        return IfNotNilInlinedNode(receiver, body_expr, source)
+
     @staticmethod
     def _try_inlining_if_else(if_true, receiver, arguments, source, mgenc):
         arg1 = arguments[0]
@@ -263,6 +286,25 @@ def _try_inlining_if_else(if_true, receiver, arguments, source, mgenc):
         false_expr = arg2.get_method().inline(mgenc)
         return IfElseInlinedNode(receiver, true_expr, false_expr, if_true, source)
 
+    @staticmethod
+    def _try_inlining_if_nil_not_nil(is_if_nil, receiver, arguments, source, mgenc):
+        arg1 = arguments[0]
+        if not isinstance(arg1, BlockNode):
+            return None
+
+        arg2 = arguments[1]
+        if not isinstance(arg2, BlockNode):
+            return None
+
+        arg1_expr = arg1.get_method().inline(mgenc)
+        arg2_expr = arg2.get_method().inline(mgenc)
+        return IfNilNotNilInlinedNode(
+            receiver,
+            arg1_expr if is_if_nil else arg2_expr,
+            arg2_expr if is_if_nil else arg1_expr,
+            source,
+        )
+
     @staticmethod
     def _try_inlining_while(while_true, receiver, arguments, source, mgenc):
         if not isinstance(receiver, BlockNode):
@@ -321,6 +363,18 @@ def _keyword_message(self, mgenc, receiver):
                     )
                     if inlined is not None:
                         return inlined
+                elif keyword == "ifNil:":
+                    inlined = self._try_inlining_if_nil(
+                        receiver, arguments, source, mgenc
+                    )
+                    if inlined is not None:
+                        return inlined
+                elif keyword == "ifNotNil:":
+                    inlined = self._try_inlining_if_not_nil(
+                        receiver, arguments, source, mgenc
+                    )
+                    if inlined is not None:
+                        return inlined
                 elif keyword == "whileTrue:":
                     inlined = self._try_inlining_while(
                         True, receiver, arguments, source, mgenc
@@ -358,6 +412,18 @@ def _keyword_message(self, mgenc, receiver):
                     )
                     if inlined is not None:
                         return inlined
+                elif keyword == "ifNil:ifNotNil:":
+                    inlined = self._try_inlining_if_nil_not_nil(
+                        True, receiver, arguments, source, mgenc
+                    )
+                    if inlined is not None:
+                        return inlined
+                elif keyword == "ifNotNil:ifNil:":
+                    inlined = self._try_inlining_if_nil_not_nil(
+                        False, receiver, arguments, source, mgenc
+                    )
+                    if inlined is not None:
+                        return inlined
 
         selector = symbol_for(keyword)
 

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)

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]

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

src/som/interpreter/ast/nodes/specialized/literal_if.py

@@ -31,6 +31,46 @@ def execute(self, frame):
         )
 
 
+class IfNilInlinedNode(ExpressionNode):
+    _immutable_fields_ = [
+        "_condition_expr?",
+        "_body_expr?",
+        "universe",
+    ]
+    _child_nodes_ = ["_condition_expr", "_body_expr"]
+
+    def __init__(self, condition_expr, body_expr, source_section):
+        ExpressionNode.__init__(self, source_section)
+        self._condition_expr = self.adopt_child(condition_expr)
+        self._body_expr = self.adopt_child(body_expr)
+
+    def execute(self, frame):
+        result = self._condition_expr.execute(frame)
+        if result is nilObject:
+            return self._body_expr.execute(frame)
+        return result
+
+
+class IfNotNilInlinedNode(ExpressionNode):
+    _immutable_fields_ = [
+        "_condition_expr?",
+        "_body_expr?",
+        "universe",
+    ]
+    _child_nodes_ = ["_condition_expr", "_body_expr"]
+
+    def __init__(self, condition_expr, body_expr, source_section):
+        ExpressionNode.__init__(self, source_section)
+        self._condition_expr = self.adopt_child(condition_expr)
+        self._body_expr = self.adopt_child(body_expr)
+
+    def execute(self, frame):
+        result = self._condition_expr.execute(frame)
+        if result is nilObject:
+            return result
+        return self._body_expr.execute(frame)
+
+
 class IfElseInlinedNode(ExpressionNode):
     _immutable_fields_ = [
         "_condition_expr?",
@@ -61,3 +101,25 @@ def execute(self, frame):
             "Would need to generalize, but we haven't implemented that "
             + "for the bytecode interpreter either"
         )
+
+
+class IfNilNotNilInlinedNode(ExpressionNode):
+    _immutable_fields_ = [
+        "_condition_expr?",
+        "_nil_expr?",
+        "_not_nil_expr?",
+        "universe",
+    ]
+    _child_nodes_ = ["_condition_expr", "_nil_expr", "_not_nil_expr"]
+
+    def __init__(self, condition_expr, nil_expr, not_nil_expr, source_section):
+        ExpressionNode.__init__(self, source_section)
+        self._condition_expr = self.adopt_child(condition_expr)
+        self._nil_expr = self.adopt_child(nil_expr)
+        self._not_nil_expr = self.adopt_child(not_nil_expr)
+
+    def execute(self, frame):
+        result = self._condition_expr.execute(frame)
+        if result is nilObject:
+            return self._nil_expr.execute(frame)
+        return self._not_nil_expr.execute(frame)