feat: check for invalid returns inside of knots

src/AstGen.zig

@@ -327,9 +327,17 @@ const GenIr = struct {
         } });
     }
 
-    fn addUnaryNode(gi: *GenIr, tag: Ir.Inst.Tag, arg: Ir.Inst.Ref) !Ir.Inst.Ref {
+    fn addUnaryNode(
+        gi: *GenIr,
+        tag: Ir.Inst.Tag,
+        arg: Ir.Inst.Ref,
+        node: *const Ast.Node,
+    ) !Ir.Inst.Ref {
         return add(gi, .{ .tag = tag, .data = .{
-            .un = .{ .lhs = arg },
+            .un = .{
+                .lhs = arg,
+                .src_offset = @intCast(node.loc.start),
+            },
         } });
     }
 
@@ -642,21 +650,21 @@ fn setCondBrPayload(
 fn unaryOp(
     gi: *GenIr,
     scope: *Scope,
-    expr_node: *const Ast.Node,
+    node: *const Ast.Node,
     op: Ir.Inst.Tag,
 ) InnerError!Ir.Inst.Ref {
-    const data = expr_node.data.bin;
+    const data = node.data.bin;
     const lhs = try expr(gi, scope, data.lhs.?);
-    return gi.addUnaryNode(op, lhs);
+    return gi.addUnaryNode(op, lhs, node);
 }
 
 fn binaryOp(
     gi: *GenIr,
     scope: *Scope,
-    expr_node: *const Ast.Node,
+    node: *const Ast.Node,
     op: Ir.Inst.Tag,
 ) InnerError!Ir.Inst.Ref {
-    const data = expr_node.data.bin;
+    const data = node.data.bin;
     const lhs = try expr(gi, scope, data.lhs.?);
     const rhs = try expr(gi, scope, data.rhs.?);
     return gi.addBinaryNode(op, lhs, rhs);
@@ -771,7 +779,7 @@ fn identifier(
 ) InnerError!Ir.Inst.Ref {
     const str = try block.astgen.strFromNode(node);
     if (scope.lookup(str.index)) |decl| {
-        return block.addUnaryNode(.load, decl.inst_index.toRef());
+        return block.addUnaryNode(.load, decl.inst_index.toRef(), node);
     }
     return block.addStrTok(.decl_ref, str.index, node.loc.start);
 }
@@ -1095,17 +1103,17 @@ fn content(
     var suppress: bool = false;
     const data = node.data.content;
     if (data.leading_glue) {
-        _ = try block.addUnaryNode(.content_glue, .none);
+        _ = try block.addUnaryNode(.content_glue, .none, node);
     }
     for (data.items) |child_node| {
         switch (child_node.tag) {
             .string_literal => {
                 const result = try stringLiteral(block, child_node);
-                _ = try block.addUnaryNode(.content_push, result);
+                _ = try block.addUnaryNode(.content_push, result, child_node);
             },
             .inline_logic_expr => {
                 const result = try inlineLogicExpr(block, scope, child_node);
-                _ = try block.addUnaryNode(.content_push, result);
+                _ = try block.addUnaryNode(.content_push, result, child_node);
             },
             .if_stmt => _ = {
                 _ = try ifStmt(block, scope, child_node);
@@ -1125,12 +1133,12 @@ fn content(
         }
     }
     if (data.trailing_glue or data.trailing_divert != null) {
-        _ = try block.addUnaryNode(.content_glue, .none);
+        _ = try block.addUnaryNode(.content_glue, .none, node);
     }
     if (!ignore_divert) {
         if (data.trailing_divert) |trailing| {
             _ = try divertExpr(block, scope, trailing);
-            _ = try block.addUnaryNode(.content_glue, .none);
+            _ = try block.addUnaryNode(.content_glue, .none, node);
         }
     }
     return suppress;
@@ -1143,7 +1151,7 @@ fn contentStmt(gi: *GenIr, scope: *Scope, node: *const Ast.Node) !void {
         supress = try content(gi, scope, item_node, false);
     }
     if (!supress) {
-        _ = try gi.addUnaryNode(.content_line, .void);
+        _ = try gi.addUnaryNode(.content_line, .void, node);
     }
 }
 
@@ -1175,7 +1183,7 @@ fn assignOp(
 
     const lhs = if (scope.lookup(name_str.index)) |decl| blk: {
         const inst_ref = decl.inst_index.toRef();
-        _ = try gi.addUnaryNode(.load, inst_ref);
+        _ = try gi.addUnaryNode(.load, inst_ref, node);
         break :blk inst_ref;
     } else blk: {
         break :blk try gi.addStrTok(.decl_ref, name_str.index, node.loc.start);
@@ -1250,7 +1258,7 @@ fn choiceStmt(gi: *GenIr, scope: *Scope, node: *const Ast.Node) InnerError!void
             _ = try blockStmt(&body_block, scope, branch_body);
         }
         if (!body_block.endsWithNoReturn()) {
-            _ = try body_block.addUnaryNode(.implicit_ret, .none);
+            _ = try body_block.addUnaryNode(.implicit_ret, .none, node);
         }
 
         const lhs_body = block_1.instructionsSliceUpto(&block_2);
@@ -1438,10 +1446,10 @@ fn divertExpr(gi: *GenIr, scope: *Scope, node: *const Ast.Node) !void {
             // TODO: Revisit this
             const str_slice = gi.astgen.tree.nodeSlice(lhs);
             if (std.mem.eql(u8, str_slice, "DONE")) {
-                _ = try gi.addUnaryNode(.done, .none);
+                _ = try gi.addUnaryNode(.done, .none, node);
                 return;
             } else if (std.mem.eql(u8, str_slice, "END")) {
-                _ = try gi.addUnaryNode(.exit, .none);
+                _ = try gi.addUnaryNode(.exit, .none, node);
                 return;
             }
             const callee = try calleeExpr(gi, scope, lhs);
@@ -1496,7 +1504,7 @@ fn returnStmt(gi: *GenIr, scope: *Scope, node: *const Ast.Node) !void {
         const arg_inst = try expr(gi, scope, lhs);
         break :blk arg_inst;
     } else .void;
-    _ = try gi.addUnaryNode(.ret, ret_arg);
+    _ = try gi.addUnaryNode(.ret, ret_arg, node);
 }
 
 fn tempDecl(gi: *GenIr, scope: *Scope, decl_node: *const Ast.Node) !void {
@@ -1588,7 +1596,7 @@ fn defaultBlock(
     const knot_inst = try decl_scope.makePayloadNode(.decl_knot);
     try blockInner(&decl_scope, scope, data.items);
     if (!decl_scope.endsWithNoReturn()) {
-        _ = try decl_scope.addUnaryNode(.implicit_ret, .none);
+        _ = try decl_scope.addUnaryNode(.implicit_ret, .none, body_node);
     }
 
     var stub_scope = decl_scope.makeSubBlock();
@@ -1634,7 +1642,8 @@ fn prototypeAndBody(
         try blockStmt(gi, scope, body);
     }
     if (!gi.endsWithNoReturn()) {
-        _ = try gi.addUnaryNode(.implicit_ret, .none);
+        // FIXME: Using `prototype_node` might break things.
+        _ = try gi.addUnaryNode(.implicit_ret, .none, prototype_node);
     }
     return .{ .decl_name = decl_name };
 }

src/Ir.zig

@@ -245,6 +245,7 @@ pub const Inst = struct {
         },
         un: struct {
             lhs: Ref,
+            src_offset: u32,
         },
         bin: struct {
             lhs: Ref,

src/Sema.zig

@@ -968,12 +968,18 @@ fn irChoiceBr(
 fn irRet(sema: *Sema, builder: *Builder, inst: Ir.Inst.Index) InnerError!void {
     const data = sema.ir.instructions[@intFromEnum(inst)].data.un;
     const lhs = sema.resolveInst(data.lhs);
-    if (lhs != .none) {
-        try builder.materialize(lhs);
-    } else {
-        try builder.addByteOp(.stream_glue);
+    const return_src: SrcLoc = .{ .src_offset = data.src_offset };
+    switch (builder.knot_tag) {
+        .function => {
+            if (lhs != .none) {
+                try builder.materialize(lhs);
+            } else {
+                try builder.addByteOp(.stream_glue);
+            }
+            try builder.addByteOp(.ret);
+        },
+        .knot => return sema.fail(return_src, "cannot return within a knot", .{}),
     }
-    try builder.addByteOp(.ret);
 }
 
 fn irImplicitRet(_: *Sema, builder: *Builder, _: Ir.Inst.Index) InnerError!void {

src/error_tests.zig

@@ -52,6 +52,19 @@ test "compiler: unknown global variable" {
     );
 }
 
+test "compiler: invalid return" {
+    try testEqual(
+        \\=== knot ===
+        \\~ return 123
+    ,
+        \\<STDIN>:2:3: error: cannot return within a knot
+        \\2    | ~ return 123
+        \\     |   ^
+        \\
+        ,
+    );
+}
+
 test "compiler: duplicate variable declarations" {
     try testEqual(
         \\VAR a = 0