Special handling of last signature in a block (#11289)

engine/runtime-compiler/src/main/scala/org/enso/compiler/pass/analyse/FramePointerAnalysis.scala

@@ -12,7 +12,8 @@ import org.enso.compiler.core.ir.{
   Function,
   Module,
   Name,
-  Pattern
+  Pattern,
+  Type
 }
 import org.enso.compiler.core.ir.module.scope.Definition
 import org.enso.compiler.core.ir.module.scope.definition.Method
@@ -163,6 +164,9 @@ case object FramePointerAnalysis extends IRPass {
         caseExpr.branches.foreach { branch =>
           processCaseBranch(branch)
         }
+      case asc: Type.Ascription =>
+        processExpression(asc.typed, graph)
+        processExpression(asc.signature, graph)
       case _ => ()
     }
     if (updateSymbols) {

engine/runtime-compiler/src/main/scala/org/enso/compiler/pass/resolve/TypeSignatures.scala

@@ -315,7 +315,15 @@ case object TypeSignatures extends IRPass {
         lastSignature = None
         res
       case a => Some(resolveExpression(a))
-    } ::: lastSignature.map(errors.Unexpected.TypeSignature(_)).toList
+    } ::: lastSignature
+      .map({
+        case asc @ Type.Ascription(_: Name, sig, comment, _, _) =>
+          asc.updateMetadata(
+            new MetadataPair(this, Signature(sig, comment))
+          )
+        case any => errors.Unexpected.TypeSignature(any)
+      })
+      .toList
 
     block.copy(
       expressions = newExpressions.init,

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

@@ -973,6 +973,59 @@ public class SignatureTest extends ContextTest {
     }
   }
 
+  @Test
+  public void returnTypeNoCheckJustSignature() throws Exception {
+    final URI uri = new URI("memory://rts.enso");
+    final Source src =
+        Source.newBuilder(
+                "enso",
+                """
+                type Integer
+
+                plusChecked a b =
+                    x:Integer
+                    x = a + b
+                    x
+                """,
+                uri.getAuthority())
+            .uri(uri)
+            .buildLiteral();
+
+    var module = ctx.eval(src);
+    var plusChecked = module.invokeMember(MethodNames.Module.EVAL_EXPRESSION, "plusChecked");
+    assertEquals(5, plusChecked.execute(2, 3).asInt());
+    var res = plusChecked.execute("a", "b");
+    assertEquals("ab", res.asString());
+  }
+
+  @Test
+  public void returnTypeCheckByLastStatement() throws Exception {
+    final URI uri = new URI("memory://rts.enso");
+    final Source src =
+        Source.newBuilder(
+                "enso",
+                """
+                from Standard.Base import Integer
+
+                plusChecked a b =
+                    x = a + b
+                    x:Integer
+                """,
+                uri.getAuthority())
+            .uri(uri)
+            .buildLiteral();
+
+    var module = ctx.eval(src);
+    var plusChecked = module.invokeMember(MethodNames.Module.EVAL_EXPRESSION, "plusChecked");
+    assertEquals(5, plusChecked.execute(2, 3).asInt());
+    try {
+      var res = plusChecked.execute("a", "b");
+      fail("Expecting an exception, not: " + res);
+    } catch (PolyglotException e) {
+      assertContains("expected expression to be Integer, but got Text", e.getMessage());
+    }
+  }
+
   /**
    * Similar scenario to {@code returnTypeCheckOptInError}, but with the opt out signature the check
    * is not currently performed.

engine/runtime-integration-tests/src/test/scala/org/enso/compiler/test/pass/resolve/TypeSignaturesTest.scala

@@ -6,6 +6,7 @@ import org.enso.compiler.core.Implicits.AsMetadata
 import org.enso.compiler.core.ir.Expression
 import org.enso.compiler.core.ir.Function
 import org.enso.compiler.core.ir.Module
+import org.enso.compiler.core.ir.Type
 import org.enso.compiler.core.ir.expression.Application
 import org.enso.compiler.core.ir.expression.errors
 import org.enso.compiler.core.ir.module.scope.Definition
@@ -236,8 +237,8 @@ class TypeSignaturesTest extends CompilerTest {
       head.getMetadata(DocumentationComments) shouldBe defined
     }
 
-    "raise an error if a signature is divorced from its definition" in {
-      block.returnValue shouldBe an[errors.Unexpected.TypeSignature]
+    "last line of a block is Type.Ascription" in {
+      block.returnValue shouldBe an[Type.Ascription]
     }
 
     "work recursively" in {

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

@@ -878,7 +878,8 @@ class IrToTruffle(
       // Type contexts aren't currently really used. But we should still check the base type.
       extractAscribedType(comment, typeInContext.typed)
     case t => {
-      t.getMetadata(TypeNames) match {
+      val res = t.getMetadata(TypeNames)
+      res match {
         case Some(
               BindingsMap
                 .Resolution(binding @ BindingsMap.ResolvedType(_, _))
@@ -1306,6 +1307,15 @@ class IrToTruffle(
         case binding: Expression.Binding => processBinding(binding)
         case caseExpr: Case =>
           processCase(caseExpr, subjectToInstrumentation)
+        case asc: Tpe.Ascription =>
+          val checkNode =
+            extractAscribedType(asc.comment.orNull, asc.signature)
+          if (checkNode != null) {
+            val body = run(asc.typed, binding, subjectToInstrumentation)
+            ReadArgumentCheckNode.wrap(body, checkNode)
+          } else {
+            processType(asc)
+          }
         case typ: Tpe => processType(typ)
         case _: Empty =>
           processEmpty()