Pattern match does not evaluate suspended atom fields

engine/runtime-integration-tests/src/test/java/org/enso/interpreter/test/LazyAtomFieldTest.java

@@ -1,5 +1,9 @@
 package org.enso.interpreter.test;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.not;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -175,6 +179,63 @@ public void toTextOnAtomWithLazyField() throws URISyntaxException {
     assertTrue(res.isString());
   }
 
+  @Test
+  public void lazyAtomFieldIsNotEvaluated_InStructuralPatternMatch() {
+    var code =
+        """
+        from Standard.Base import Nothing, IO
+
+        type My_Ref
+            Lazy ~lazy
+            Eager eager
+
+        main =
+            v1 = My_Ref.Lazy <|
+                IO.println "Computing v1"
+                42
+
+            case v1 of
+                My_Ref.Eager e -> e
+                My_Ref.Lazy _ -> Nothing
+        """;
+    var res = ctxRule.evalModule(code);
+    assertThat(res.isNull(), is(true));
+    assertThat(
+        "Lazy field is not evaluated in pattern match",
+        ctxRule.getStdOut(),
+        not(containsString("Computing v1")));
+  }
+
+  @Test
+  public void lazyAtomFieldIsNotEvaluated_InStructuralPatternMatch_WithBoundedField() {
+    var code =
+        """
+        from Standard.Base import Nothing, IO, False
+
+        type My_Ref
+            Lazy ~lazy
+            Eager eager
+
+        main =
+            v2 = My_Ref.Lazy <|
+                IO.println "Computing v2"
+                42
+
+            should_compute = False
+
+            case v2 of
+                My_Ref.Eager e -> e
+                My_Ref.Lazy l ->
+                    if should_compute then l else Nothing
+        """;
+    var res = ctxRule.evalModule(code);
+    assertThat(res.isNull(), is(true));
+    assertThat(
+        "Lazy field is not evaluated in pattern match with bounded field",
+        ctxRule.getStdOut(),
+        not(containsString("Computing v2")));
+  }
+
   private void checkNumHolder(String typeDefinition) throws Exception {
     var code =
         "polyglot java import java.util.Random as R\n"

engine/runtime/src/main/java/org/enso/interpreter/node/controlflow/caseexpr/ConstructorBranchNode.java

@@ -1,6 +1,7 @@
 package org.enso.interpreter.node.controlflow.caseexpr;
 
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.RootCallTarget;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Cached.Shared;
@@ -10,8 +11,11 @@
 import com.oracle.truffle.api.library.CachedLibrary;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.nodes.NodeInfo;
+import com.oracle.truffle.api.nodes.RootNode;
 import com.oracle.truffle.api.profiles.CountingConditionProfile;
+import org.enso.interpreter.EnsoLanguage;
 import org.enso.interpreter.node.expression.builtin.meta.IsValueOfTypeNode;
+import org.enso.interpreter.runtime.callable.function.Function;
 import org.enso.interpreter.runtime.data.EnsoMultiValue;
 import org.enso.interpreter.runtime.data.atom.Atom;
 import org.enso.interpreter.runtime.data.atom.AtomConstructor;
@@ -85,9 +89,41 @@ private static Object[] fieldsFromObject(
       Object obj, AtomConstructor cons, StructsLibrary structsLib) {
     CompilerAsserts.partialEvaluationConstant(cons);
     var arr = new Object[cons.getArity()];
+    var fields = cons.getFields();
     for (var i = 0; i < arr.length; i++) {
-      arr[i] = structsLib.getField(obj, i);
+      // Wrap suspended fields in functions, so that they are not evaluated
+      // too eagerly.
+      if (fields[i].isSuspended() && obj instanceof Atom atom) {
+        var func = createSuspendedFieldFunction(atom, i);
+        arr[i] = func;
+      } else {
+        arr[i] = structsLib.getField(obj, i);
+      }
     }
     return arr;
   }
+
+  @TruffleBoundary
+  private static Function createSuspendedFieldFunction(Atom atom, int fieldIdx) {
+    var getFieldNode = new GetSuspendedFieldNode(atom, fieldIdx);
+    return Function.fullyApplied(getFieldNode.getCallTarget());
+  }
+
+  private static final class GetSuspendedFieldNode extends RootNode {
+    private final Atom atom;
+    private final int fieldIdx;
+    @Child private StructsLibrary structsLib = StructsLibrary.getFactory().createDispatched(3);
+
+    GetSuspendedFieldNode(Atom atom, int fieldIdx) {
+      super(EnsoLanguage.get(null));
+      this.atom = atom;
+      this.fieldIdx = fieldIdx;
+    }
+
+    @Override
+    public Object execute(VirtualFrame frame) {
+      // Forces evaluation of the suspended field.
+      return structsLib.getField(atom, fieldIdx);
+    }
+  }
 }

engine/runtime/src/main/scala/org/enso/interpreter/runtime/IrToTruffle.scala

@@ -1481,17 +1481,6 @@ private[runtime] class IrToTruffle(
             )
           }
 
-          val fieldNames   = cons.unsafeFieldsAsNamed
-          val fieldsAsArgs = fieldNames.map(genArgFromMatchField)
-
-          val branchCodeNode = childProcessor.processFunctionBody(
-            fieldsAsArgs,
-            branch.expression,
-            branch.location,
-            subjectToInstrumentation = subjectToInstrumentation,
-            defineRoot               = false
-          )
-
           constructor match {
             case err: errors.Resolution =>
               Left(BadPatternMatch.NonVisibleConstructor(err.name))
@@ -1502,6 +1491,14 @@ private[runtime] class IrToTruffle(
                 case Some(
                       BindingsMap.Resolution(BindingsMap.ResolvedModule(mod))
                     ) =>
+                  val branchCodeNode =
+                    createBranchCodeNodeForConstructorPattern(
+                      branch,
+                      cons,
+                      None,
+                      childProcessor,
+                      subjectToInstrumentation
+                    )
                   Right(
                     ObjectEqualityBranchNode.build(
                       branchCodeNode.getCallTarget,
@@ -1511,11 +1508,19 @@ private[runtime] class IrToTruffle(
                   )
                 case Some(
                       BindingsMap.Resolution(
-                        BindingsMap.ResolvedConstructor(tp, cons)
+                        BindingsMap.ResolvedConstructor(tp, resolvedCons)
                       )
                     ) =>
                   val atomCons =
-                    asType(tp).getConstructors.get(cons.name)
+                    asType(tp).getConstructors.get(resolvedCons.name)
+                  val branchCodeNode =
+                    createBranchCodeNodeForConstructorPattern(
+                      branch,
+                      cons,
+                      Some(atomCons),
+                      childProcessor,
+                      subjectToInstrumentation
+                    )
                   val r = if (atomCons == getBuiltins.bool().getTrue) {
                     BooleanBranchNode.build(
                       true,
@@ -1543,6 +1548,14 @@ private[runtime] class IrToTruffle(
                     ) =>
                   val tpe =
                     asType(binding)
+                  val branchCodeNode =
+                    createBranchCodeNodeForConstructorPattern(
+                      branch,
+                      cons,
+                      None,
+                      childProcessor,
+                      subjectToInstrumentation
+                    )
                   val polyglot = getBuiltins.polyglot
                   val branchNode = if (tpe == polyglot) {
                     PolyglotBranchNode.build(
@@ -1563,6 +1576,14 @@ private[runtime] class IrToTruffle(
                         BindingsMap.ResolvedPolyglotSymbol(mod, symbol)
                       )
                     ) =>
+                  val branchCodeNode =
+                    createBranchCodeNodeForConstructorPattern(
+                      branch,
+                      cons,
+                      None,
+                      childProcessor,
+                      subjectToInstrumentation
+                    )
                   val polyglotSymbol =
                     asScope(mod.unsafeAsModule())
                       .getPolyglotSymbolSupplier(symbol.name)
@@ -1581,6 +1602,14 @@ private[runtime] class IrToTruffle(
                         BindingsMap.ResolvedPolyglotField(typ, symbol)
                       )
                     ) =>
+                  val branchCodeNode =
+                    createBranchCodeNodeForConstructorPattern(
+                      branch,
+                      cons,
+                      None,
+                      childProcessor,
+                      subjectToInstrumentation
+                    )
                   val mod = typ.module
                   val polyClass = asScope(mod.unsafeAsModule())
                     .getPolyglotSymbolSupplier(typ.symbol.name)
@@ -1801,6 +1830,62 @@ private[runtime] class IrToTruffle(
       }
     }
 
+    /** Case branch of a [[Pattern.Constructor]] is represented as a function with parameters matching the
+      * fields of the constructor. This method converts those fields into function arguments.
+      * If a field is suspended, the corresponding argument is marked as suspended too.
+      * @param cons Constructor pattern.
+      * @param resolvedCons The actual constructor resolved at runtime.
+      * @return
+      */
+    private def atomFieldsAsArguments(
+      cons: Pattern.Constructor,
+      resolvedCons: AtomConstructor
+    ): List[DefinitionArgument] = {
+      val resolvedFields = resolvedCons.getFields.toList
+      val fieldNames     = cons.unsafeFieldsAsNamed
+      fieldNames.zip(resolvedFields).map { case (fieldName, resolvedField) =>
+        DefinitionArgument.Specified
+          .builder()
+          .name(fieldName.name)
+          .suspended(resolvedField.isSuspended)
+          .location(fieldName.identifiedLocation)
+          .passData(fieldName.name.passData())
+          .diagnostics(fieldName.name.diagnostics())
+          .build()
+      }
+    }
+
+    private def createBranchCodeNodeForConstructorPattern(
+      branch: Case.Branch,
+      cons: Pattern.Constructor,
+      resolvedConsOpt: Option[AtomConstructor],
+      childProcessor: ExpressionProcessor,
+      subjectToInstrumentation: Boolean
+    ): CreateFunctionNode = {
+      val fieldNames = cons.unsafeFieldsAsNamed
+      resolvedConsOpt match {
+        case None =>
+          val fieldsAsArgs = fieldNames.map(genArgFromMatchField)
+          childProcessor.processFunctionBody(
+            fieldsAsArgs,
+            branch.expression,
+            branch.location,
+            subjectToInstrumentation = subjectToInstrumentation,
+            defineRoot               = false
+          )
+        case Some(resolvedOpt) =>
+          val fieldsAsSuspendedArgs =
+            atomFieldsAsArguments(cons, resolvedOpt)
+          childProcessor.processFunctionBody(
+            fieldsAsSuspendedArgs,
+            branch.expression,
+            branch.location,
+            subjectToInstrumentation = subjectToInstrumentation,
+            defineRoot               = false
+          )
+      }
+    }
+
     /* Note [Pattern Match Fallbacks]
      * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      * Enso in its current state has no coverage checking for constructors on

test/Base_Tests/src/Runtime/Lazy_Spec.enso

@@ -14,6 +14,12 @@ type Lazy
 type Generator
     Value n ~next
 
+type Wrapper
+    Value field
+
+type Lazy_Fields
+    Value ~x ~y ~z
+
 
 add_specs suite_builder = suite_builder.group "Lazy" group_builder->
     group_builder.specify "should compute the result only once" <|
@@ -110,6 +116,77 @@ add_specs suite_builder = suite_builder.group "Lazy" group_builder->
         atom.get . should_equal 2 # Force evaluation
         atom.to_text . should_equal "(Lazy.Value 2)"
 
+    group_builder.specify "lazy atom field is not evaluated in structural pattern match" <|
+        ref = Ref.new 0
+        l = Lazy.new <|
+            ref.modify (_ + 1)
+        case l of
+            Lazy.Value _ -> Nothing
+            _ -> Test.fail "should match Lazy.Value ctor"
+        ref.get . should_equal 0
+
+    group_builder.specify "lazy atom field is not evaluated in structural pattern match with bounded field" <|
+        ref = Ref.new 0
+        l = Lazy.new <|
+            ref.modify (_ + 1)
+        should_compute = False
+        case l of
+            Lazy.Value x ->
+                if should_compute then x else Nothing
+            _ -> Test.fail "should match Lazy.Value ctor"
+        ref.get . should_equal 0
+
+    group_builder.specify "lazy atom field is evaluated just once in structural pattern match" <|
+        ref = Ref.new 0
+        l = Lazy.new <|
+            ref.modify (_ + 1)
+        should_compute = True
+        case l of
+            Lazy.Value x ->
+                ref.get . should_equal 0
+                if should_compute then x else Nothing
+            _ -> Test.fail "should match Lazy.Value ctor"
+        ref.get . should_equal 1
+
+    group_builder.specify "only certain lazy fields are forced in structural pattern match" <|
+        ref = Ref.new 0
+        computation =
+            ref.modify (_ + 1)
+        l = Lazy_Fields.Value computation computation computation
+        res = case l of
+            Lazy_Fields.Value x _ z ->
+                [x, z]
+            _ -> Test.fail "should match Lazy_Fields.Value ctor"
+        # Ref.modify returns the old value
+        res . should_equal [0, 1]
+        ref.get . should_equal 2
+
+    group_builder.specify "do not recompute lazy atom field after structural pattern match" <|
+        ref = Ref.new 0
+        l = Lazy.new <|
+            ref.modify (_ + 1)
+        case l of
+            Lazy.Value x ->
+                x . should_equal 0
+            _ -> Test.fail "should match Lazy.Value ctor"
+        case l of
+            Lazy.Value x ->
+                x . should_equal 0
+            _ -> Test.fail "should match Lazy.Value ctor"
+        ref.get . should_equal 1
+
+
+    group_builder.specify "lazy atom field is not evaluated in structural pattern match when wrapped" <|
+        ref = Ref.new 0
+        l = Lazy.new <|
+            ref.modify (_ + 1)
+        wrap = Wrapper.Value l
+        case wrap of
+            Wrapper.Value (Lazy.Value _) -> Nothing
+            _ -> Test.fail "should match Wrapper.Value ctor"
+        ref.get . should_equal 0
+
+
 main filter=Nothing =
     suite = Test.build suite_builder->
         add_specs suite_builder