Added support for inlining #ifNil:* and #ifNotNil:* in the AST interpreter

Signed-off-by: Stefan Marr <[email protected]>

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/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)

tests/test_ast_inlining.py

@@ -22,6 +22,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.interpreter.ast.nodes.variable_node import (
@@ -205,6 +208,29 @@ def test_if_arg(mgenc, if_selector, expected_bool, unexpected_bool):
     assert if_node._not_expected_bool is unexpected_bool
 
 
+@pytest.mark.parametrize(
+    "if_selector,expected_class",
+    [
+        ("ifNil:", IfNilInlinedNode),
+        ("ifNotNil:", IfNotNilInlinedNode),
+    ],
+)
+def test_if_nil_arg(mgenc, if_selector, expected_class):
+    ast = parse_method(
+        mgenc,
+        """
+        test: arg = (
+            #start.
+            self method IF_SELECTOR [ arg ].
+            #end
+        )""".replace(
+            "IF_SELECTOR", if_selector
+        ),
+    )
+
+    assert isinstance(ast._exprs[1], expected_class)
+
+
 def test_if_true_and_inc_field(cgenc, mgenc):
     add_field(cgenc, "field")
     ast = parse_method(
@@ -430,6 +456,35 @@ def test_if_true_if_false_return(mgenc, sel1, sel2, expected_bool, unexpected_bo
     assert if_node._not_expected_bool is unexpected_bool
 
 
+@pytest.mark.parametrize(
+    "sel1,sel2,expected_nil_expr,expected_not_nil_expr",
+    [
+        ("ifNil:", "ifNotNil:", ReturnLocalNode, LiteralNode),
+        ("ifNotNil:", "ifNil:", LiteralNode, ReturnLocalNode),
+    ],
+)
+def test_if_nil_not_nil_return(
+    mgenc, sel1, sel2, expected_nil_expr, expected_not_nil_expr
+):
+    seq = parse_method(
+        mgenc,
+        """
+        test: arg1 with: arg2 = (
+            #start.
+            ^ self method SEL1 [ ^ arg1 ] SEL2 [ #foo ]
+        )""".replace(
+            "SEL1", sel1
+        ).replace(
+            "SEL2", sel2
+        ),
+    )
+
+    assert isinstance(seq._exprs[1], IfNilNotNilInlinedNode)
+    if_node = seq._exprs[1]
+    assert isinstance(if_node._nil_expr, expected_nil_expr)
+    assert isinstance(if_node._not_nil_expr, expected_not_nil_expr)
+
+
 @pytest.mark.parametrize(
     "while_sel,expected_bool,unexpected_bool",
     [