[Heavy] Add syntax-source-loc for propagating source locations of syntax-rules pattern variables

ricejasonf · ricejasonf · commit ced358645f58 · 2025-10-25T15:51:41.000-07:00
diff --git a/heavy/include/heavy/Builtins.h b/heavy/include/heavy/Builtins.h
@@ -29,6 +29,8 @@
 #define HEAVY_BASE_VAR(NAME)          ::heavy::builtins_var::NAME
 #define HEAVY_BASE_VAR_STR(NAME)      HEAVY_BASE_VAR_STR__##NAME
 #define HEAVY_BASE_VAR_STR__error     HEAVY_BASE_LIB_STR "V5Serror"
+#define HEAVY_BASE_VAR_STR__source_loc \
+  HEAVY_BASE_LIB_STR "V6Ssourcemi3Sloc"
 
 // Forward declare vars that are used in the C++ codebase.
 namespace heavy::builtins_var {
diff --git a/heavy/include/heavy/Context.h b/heavy/include/heavy/Context.h
@@ -344,14 +344,17 @@ class Context : public ContinuationStack<Context>,
   Pair*       CreatePair(Value V1) {
     return new (*this) Pair(V1, CreateEmpty());
   }
-  Pair*       CreatePair(Value V1, Value V2, Pair* SourcePair) {
-    heavy::SourceLocation Loc = SourcePair->getSourceLocation();
+  Pair* CreatePair(Value V1, Value V2, heavy::SourceLocation MaybeLoc) {
     // Try to preserve source locations.
-    if (Loc.isValid())
-      return CreatePairWithSource(V1, V2, Loc);
+    if (MaybeLoc.isValid())
+      return CreatePairWithSource(V1, V2, MaybeLoc);
     else
       return CreatePair(V1, V2);
   }
+  Pair* CreatePair(Value V1, Value V2, Value MaybeSource) {
+    heavy::SourceLocation Loc = MaybeSource.getSourceLocation();
+    return CreatePair(V1, V2, Loc);
+  }
   PairWithSource* CreatePairWithSource(Value V1, Value V2,
                                        SourceLocation Loc) {
     return new (*this) PairWithSource(V1, V2, Loc);
@@ -419,6 +422,7 @@ class Context : public ContinuationStack<Context>,
   Value CreateSyntaxClosure(SourceLocation Loc, Value Node, Value Env);
   SyntaxClosure* CreateSyntaxClosure(SourceLocation Loc, Symbol* S,
                                      Value Env);
+  Value UnwrapSyntaxClosure(Value V);
 
   SourceValue* CreateSourceValue(SourceLocation Loc) {
     return new (*this) SourceValue(Loc);
diff --git a/heavy/include/heavy/ContinuationStack.h b/heavy/include/heavy/ContinuationStack.h
@@ -266,11 +266,8 @@ class ContinuationStack {
         break;
       }
       default:
-        String* Msg = Context.CreateString(
-          "invalid operator for call expression: ",
-          getKindName(Callee.getKind())
-        );
-        Context.RaiseError(Msg, Callee);
+        Context.RaiseError(
+            "invalid operator for call expression: {}", Callee);
       }
 
       // This means a C++ function was not written correctly.
diff --git a/heavy/include/heavy/Ops.td b/heavy/include/heavy/Ops.td
@@ -580,17 +580,18 @@ def heavy_RenameEnvOp : HeavyOp<"rename_env"> {
 def heavy_SourceLocOp : HeavyOp<"source-loc"> {
   let summary = "source-loc";
   let description = [{
-    Encode a source location that is specific
-    to an instance of a SourceManager. Serialization
-    should not save this encoding.
-    This is used to allow passing source locations of
-    syntax as values.
+    Get the source location from a HeavyValue object.
+    This is primarily used to propagate source locations
+    of the inputs to pattern variables.
   }];
 
-  let arguments = (ins);
+  let arguments = (ins HeavyValue:$arg);
   let results   = (outs HeavyValue:$result);
-  let builders  = [OpBuilder<(ins)>];
+  let builders = [
+    OpBuilder<(ins "::mlir::Value":$arg)>
+  ];
 }
+
 def heavy_SetOp : HeavyOp<"set"> {
   let summary = "set";
   let description = [{
@@ -629,10 +630,12 @@ def heavy_SyntaxClosureOp : HeavyOp<"syntax_closure"> {
     macro expansion.
   }];
 
-  let arguments = (ins HeavyValue:$input, HeavyValue:$env);
+  let arguments = (ins HeavyValue:$sourceVal, HeavyValue:$input,
+                       HeavyValue:$env);
   let results   = (outs HeavyValue:$result);
   let builders  = [
-    OpBuilder<(ins "::mlir::Value":$input, "::mlir::Value":$env)>
+    OpBuilder<(ins "::mlir::Value":$sourceVal, "::mlir::Value":$input,
+                   "::mlir::Value":$env)>
   ];
 }
 
diff --git a/heavy/include/heavy/base.sld b/heavy/include/heavy/base.sld
@@ -47,6 +47,7 @@
     car
     cdr
     cons
+    source-cons
     dump
     dynamic-wind
     eq?
diff --git a/heavy/include/heavy/base/r7rs-syntax.sld b/heavy/include/heavy/base/r7rs-syntax.sld
@@ -48,8 +48,6 @@
         ((do "step" x y)
          y)))
 
-    ; FIXME cond should be in the environment within syntax body.
-    ;       (this applies to all define-syntax)
     (define-syntax cond
       (syntax-rules (else =>)
         ((cond (else result1 result2 ...))
diff --git a/heavy/lib/Builtins.cpp b/heavy/lib/Builtins.cpp
@@ -58,6 +58,7 @@ heavy::ExternFunction is_zero;
 heavy::ExternFunction list;
 heavy::ExternFunction length;
 heavy::ExternFunction cons;
+heavy::ExternFunction source_cons;
 heavy::ExternFunction car;
 heavy::ExternFunction cdr;
 heavy::ExternFunction append;
@@ -795,6 +796,15 @@ void cons(Context& C, ValueRefs Args) {
   C.Cont(C.CreatePair(Args[0], Args[1]));
 }
 
+// (source-cons car cdr src)
+// Where `src` can be any value that might
+// have a valid source location.
+void source_cons(Context& C, ValueRefs Args) {
+  if (Args.size() != 3)
+    return C.RaiseError("invalid arity");
+  C.Cont(C.CreatePair(Args[0], Args[1], Args[2]));
+}
+
 void car(Context& C, ValueRefs Args) {
   if (Args.size() != 1)
     return C.RaiseError("invalid arity");
@@ -1263,6 +1273,7 @@ void HEAVY_BASE_INIT(heavy::Context& Context) {
   HEAVY_BASE_VAR(list)    = heavy::builtins::list;
   HEAVY_BASE_VAR(length)  = heavy::builtins::length;
   HEAVY_BASE_VAR(cons)    = heavy::builtins::cons;
+  HEAVY_BASE_VAR(source_cons) = heavy::builtins::source_cons;
   HEAVY_BASE_VAR(car)     = heavy::builtins::car;
   HEAVY_BASE_VAR(cdr)     = heavy::builtins::cdr;
   HEAVY_BASE_VAR(append)  = heavy::builtins::append;
@@ -1359,6 +1370,7 @@ void HEAVY_BASE_LOAD_MODULE(heavy::Context& Context) {
     {"list",    HEAVY_BASE_VAR(list)},
     {"length",  HEAVY_BASE_VAR(length)},
     {"cons",    HEAVY_BASE_VAR(cons)},
+    {"source-cons", HEAVY_BASE_VAR(source_cons)},
     {"car",     HEAVY_BASE_VAR(car)},
     {"cdr",     HEAVY_BASE_VAR(cdr)},
     {"append",  HEAVY_BASE_VAR(append)},
diff --git a/heavy/lib/Context.cpp b/heavy/lib/Context.cpp
@@ -368,7 +368,7 @@ class Writer : public ValueVisitor<Writer> {
 
   void VisitSourceValue(SourceValue* SV) {
     heavy::SourceLocation Loc = SV->getSourceLocation();
-    OS << "\n#<SourceValue {"
+    OS << "#<SourceValue {"
        << Loc.getEncoding()
        << "}>";
   }
@@ -1310,10 +1310,8 @@ void Context::Raise(Value Obj) {
 
 void Context::RaiseError(String* Msg, llvm::ArrayRef<Value> IrrArgs) {
   Value IrrList = Empty();
-  for (Value Irr : llvm::reverse(IrrArgs)) {
-    setLoc(Irr.getSourceLocation());
+  for (Value Irr : llvm::reverse(IrrArgs))
     IrrList = CreatePair(Irr, IrrList);
-  }
   Value Error = CreateError(this->Loc, Msg, IrrList);
   Raise(Error);
 }
@@ -1672,14 +1670,30 @@ EnvFrame* Context::CreateEnvFrame(llvm::ArrayRef<Value> Ids, bool IsLambdaScope)
 // Create syntax with Symbols rebuilt as SyntaxClosures.
 Value Context::CreateSyntaxClosure(SourceLocation Loc, Value Node,
                                                        Value Env) {
-  return RebuildLiteral(Node, Env);
+  // Substitute all nested symbols with a SyntaxClosure;
+  Value Rebuilt = RebuildLiteral(Node, Env);
+  // Prevent directly nesting syntax closures.
+  if (isa<SyntaxClosure>(Rebuilt))
+    return Rebuilt;
+
+  // Wrap in a SyntaxClosure to keep the Loc available.
+  return new (*this) SyntaxClosure(Loc, Env, Rebuilt);
 }
 
 SyntaxClosure* Context::CreateSyntaxClosure(SourceLocation Loc, Symbol* S,
                                             Value Env) {
   return new (*this) SyntaxClosure(Loc, Env, Value(S));
 }
 
+// Unwrap any SyntaxClosure that is not an Identifier.
+// (We only keep these around for the SourceLocation.)
+Value Context::UnwrapSyntaxClosure(Value V) {
+  if (auto* SC = dyn_cast<SyntaxClosure>(V);
+      SC && !isa<Symbol>(SC->Node))
+    return SC->Node;
+  return V;
+}
+
 EnvEntry Context::GetSyntax(EnvEntry Entry) {
   heavy::Value V = Entry.Value;
   if (!V)
diff --git a/heavy/lib/Dialect.cpp b/heavy/lib/Dialect.cpp
@@ -240,9 +240,10 @@ void SpliceOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
 }
 
 void SyntaxClosureOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
-                            mlir::Value Input, mlir::Value Env) {
+                            mlir::Value SourceVal, mlir::Value Input,
+                            mlir::Value Env) {
   mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
-  SyntaxClosureOp::build(B, OpState, HeavyValueT, Input, Env);
+  SyntaxClosureOp::build(B, OpState, HeavyValueT, SourceVal, Input, Env);
 }
 
 void SyntaxOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
@@ -273,9 +274,10 @@ void UndefinedOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState) {
   UndefinedOp::build(B, OpState, B.getType<HeavyValueTy>());
 }
 
-void SourceLocOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState) {
+void SourceLocOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
+                        mlir::Value Arg) {
   mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
-  SourceLocOp::build(B, OpState, HeavyValueT);
+  SourceLocOp::build(B, OpState, HeavyValueT, Arg);
 }
 
 using namespace mlir;
diff --git a/heavy/lib/OpEval.cpp b/heavy/lib/OpEval.cpp
@@ -85,6 +85,13 @@ class OpEvalImpl {
     return V;
   }
 
+  // Get a value for pattern matching
+  // that can ignore syntax closures
+  // (which does not include symbols.)
+  heavy::Value getPatternValue(mlir::Value M) {
+    return Context.UnwrapSyntaxClosure(getValue(M));
+  }
+
   static heavy::SourceLocation getSourceLocation(mlir::Location Loc) {
     return heavy::SourceLocation(mlir::OpaqueLoc
       ::getUnderlyingLocationOrNull<heavy::SourceLocationEncoding*>(Loc));
@@ -674,7 +681,7 @@ class OpEvalImpl {
 
   BlockItrTy Visit(MatchOp Op) {
     heavy::Value P = Op.getVal().getValue(Context);
-    heavy::Value E = getValue(Op.getInput());
+    heavy::Value E = getPatternValue(Op.getInput());
     heavy::EnvEntry Entry;
     if (isIdentifier(E) && isIdentifier(P)) {
       // If the symbol is in the environment we can skip
@@ -699,7 +706,7 @@ class OpEvalImpl {
   }
 
   BlockItrTy Visit(MatchPairOp Op) {
-    heavy::Value E = getValue(Op.getInput());
+    heavy::Value E = getPatternValue(Op.getInput());
     if (auto* Pair = dyn_cast<heavy::Pair>(E)) {
       setValue(Op.getCar(), Pair->Car);
       setValue(Op.getCdr(), Pair->Cdr);
@@ -712,7 +719,7 @@ class OpEvalImpl {
   BlockItrTy Visit(MatchTailOp Op) {
     // Here we include the last cdr in the length
     // of a list proper or improper.
-    heavy::Value E = getValue(Op.getInput());
+    heavy::Value E = getPatternValue(Op.getInput());
     uint32_t TargetLen = Op.getLength();
     assert(TargetLen >= 1 && "expecting positive length");
     uint32_t TotalLen = 1;
@@ -828,7 +835,7 @@ class OpEvalImpl {
         break;
       auto Scope = ValueMapScope(ValueMap);
       BlockItrTy Itr = Op.getBody().front().begin();
-      setValue(BodyArg, Pair->Car);
+      setValue(BodyArg, Pair);
       while (Itr != BlockItrTy() && Itr != BlockItrTy(Op))
         Itr = Visit(&*Itr);
       // Check subpattern failure.
@@ -846,6 +853,7 @@ class OpEvalImpl {
   }
 
   BlockItrTy Visit(ResolveOp Op) {
+    heavy::SourceLocation Loc = getSourceLocation(Op.getLoc());
     if (auto ParentOp = dyn_cast<SubpatternOp>(Op->getParentOp())) {
       // We are at the terminator of a subpattern body region.
       assert(ParentOp.getPacks().size() ==
@@ -856,7 +864,7 @@ class OpEvalImpl {
       // Construct each pack as a reversed ordered list.
       for (unsigned i = 0; i < ResolveArgs.size(); i++) {
         Value Pack = Context.CreatePair(getValue(ResolveArgs[i]),
-                                        getValue(MResults[i]));
+                                        getValue(MResults[i]), Loc);
         setValueParentScope(MResults[i], Pack);
       }
       return BlockItrTy();
@@ -866,7 +874,8 @@ class OpEvalImpl {
     assert(Op.getArgs().size() == 1 && "expecting single result");
     heavy::Value CurrentResult = getValue(Op.getArgs().front());
     heavy::Value Result = Context.CreatePair(CurrentResult,
-                                      getValue(ExpandPacksOp.getResult()));
+                                      getValue(ExpandPacksOp.getResult()),
+                                      Loc);
     setValueParentScope(ExpandPacksOp.getResult(), Result);
     return BlockItrTy();
   }
@@ -882,9 +891,11 @@ class OpEvalImpl {
 
   BlockItrTy Visit(SyntaxClosureOp Op) {
     // Create a SyntaxClosure with the provided env argument.
+    heavy::Value SourceVal = getValue(Op.getSourceVal());
+    heavy::SourceLocation Loc = SourceVal.getSourceLocation();
     heavy::Value Input = getValue(Op.getInput());
     heavy::Value Env = getValue(Op.getEnv());
-    Value SC = Context.CreateSyntaxClosure(Context.getLoc(), Input, Env);
+    Value SC = Context.CreateSyntaxClosure(Loc, Input, Env);
     setValue(Op.getResult(), SC);
 
     return next(Op);
@@ -921,7 +932,8 @@ class OpEvalImpl {
   }
 
   BlockItrTy Visit(SourceLocOp Op) {
-    heavy::SourceLocation Loc = getSourceLocation(Op.getLoc());
+    heavy::Value Arg = getValue(Op.getArg());
+    heavy::SourceLocation Loc = Arg.getSourceLocation();
     heavy::Value V = Context.CreateSourceValue(Loc);
     setValue(Op.getResult(), V);
     return next(Op);
diff --git a/heavy/lib/OpGen.cpp b/heavy/lib/OpGen.cpp
@@ -524,6 +524,7 @@ mlir::Value OpGen::createLambda(Value Formals, Value Body,
 std::pair<heavy::FuncOp, heavy::EnvFrame*>
 OpGen::createLambdaFunction(Value Formals, SourceLocation Loc,
                             llvm::StringRef Name) {
+  Formals = Context.UnwrapSyntaxClosure(Formals);
   // Flush any internal definitions from containing body if any.
   if (IsLocalDefineAllowed)
     FinishLocalDefines();
@@ -1188,7 +1189,13 @@ mlir::Value OpGen::VisitExternName(ExternName* EN) {
 }
 
 mlir::Value OpGen::VisitSyntaxClosure(SyntaxClosure* SC) {
-  return VisitSymbol(cast<Symbol>(SC->Node), SC->Env);
+  // A SyntaxClosure wraps either a Symbol or a rebuilt expression that
+  // substitutes all symbols as SyntaxClosures so we can simply unwrap
+  // the latter.
+  if (auto* S = dyn_cast<Symbol>(SC->Node))
+    return VisitSymbol(S, SC->Env);
+  else
+    return Visit(SC->Node);
 }
 
 mlir::Value OpGen::VisitSymbol(Symbol* S, Value ClosedEnv) {
diff --git a/heavy/lib/PatternTemplate.h b/heavy/lib/PatternTemplate.h
@@ -101,9 +101,17 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
       OpGen.SetError(
         "pattern variable name appears in pattern multiple times", S);
     }
+
+    // Use any containing pair to carry source location information.
+    mlir::Value SourceVal;
+    if (auto MatchPairOp = E.getDefiningOp<heavy::MatchPairOp>())
+      SourceVal = MatchPairOp.getInput();
+    else
+      SourceVal = E;
+
     // Create a local variable with a SyntaxClosure of E as the initializer.
     heavy::SourceLocation Loc = S->getSourceLocation();
-    auto SynClo = OpGen.create<SyntaxClosureOp>(Loc, E, EnvArg);
+    auto SynClo = OpGen.create<SyntaxClosureOp>(Loc, SourceVal, E, EnvArg);
     heavy::Context& C = OpGen.getContext();
     Binding* B = C.CreateBinding(S, SynClo.getOperation());
     C.PushLocalBinding(B);
@@ -152,8 +160,10 @@ class PatternTemplate : ValueVisitor<PatternTemplate, mlir::Value> {
       OpGen.Builder.setInsertionPointToStart(&Block);
       mlir::Type HeavyValueT = OpGen.Builder.getType<HeavyValueTy>();
       mlir::Location MLoc = createLoc(Loc);
+      // The block argument is the current pair from the input.
       mlir::Value BodyArg = Block.addArgument(HeavyValueT, MLoc);
-      Visit(P, BodyArg);
+      auto PairVal = OpGen.create<MatchPairOp>(Loc, BodyArg);
+      Visit(P, PairVal.getCar());
 
       // Create range for pack values by finding the
       // SyntaxClosureOps in the Body region.
diff --git a/heavy/lib/TemplateBase.h b/heavy/lib/TemplateBase.h
diff --git a/heavy/lib/TemplateGen.h b/heavy/lib/TemplateGen.h
diff --git a/heavy/test/Evaluate/source-loc.scm b/heavy/test/Evaluate/source-loc.scm

Original file line number	Diff line number	Diff line change
`@@ -266,11 +266,8 @@ class ContinuationStack {`
`266`	`266`	`break;`
`267`	`267`	`}`
`268`	`268`	`default:`
`269`		`- String* Msg = Context.CreateString(`
`270`		`- "invalid operator for call expression: ",`
`271`		`- getKindName(Callee.getKind())`
`272`		`- );`
`273`		`- Context.RaiseError(Msg, Callee);`
	`269`	`+ Context.RaiseError(`
	`270`	`+ "invalid operator for call expression: {}", Callee);`
`274`	`271`	`}`
`275`	`272`
`276`	`273`	`// This means a C++ function was not written correctly.`
Original file line number	Diff line number	Diff line change
`@@ -240,9 +240,10 @@ void SpliceOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,`
`240`	`240`	`}`
`241`	`241`
`242`	`242`	`void SyntaxClosureOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,`
`243`		`- mlir::Value Input, mlir::Value Env) {`
	`243`	`+ mlir::Value SourceVal, mlir::Value Input,`
	`244`	`+ mlir::Value Env) {`
`244`	`245`	`mlir::Type HeavyValueT = B.getType<HeavyValueTy>();`
`245`		`- SyntaxClosureOp::build(B, OpState, HeavyValueT, Input, Env);`
	`246`	`+ SyntaxClosureOp::build(B, OpState, HeavyValueT, SourceVal, Input, Env);`
`246`	`247`	`}`
`247`	`248`
`248`	`249`	`void SyntaxOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,`
`@@ -273,9 +274,10 @@ void UndefinedOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState) {`
`273`	`274`	`UndefinedOp::build(B, OpState, B.getType<HeavyValueTy>());`
`274`	`275`	`}`
`275`	`276`
`276`		`-void SourceLocOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState) {`
	`277`	`+void SourceLocOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,`
	`278`	`+ mlir::Value Arg) {`
`277`	`279`	`mlir::Type HeavyValueT = B.getType<HeavyValueTy>();`
`278`		`- SourceLocOp::build(B, OpState, HeavyValueT);`
	`280`	`+ SourceLocOp::build(B, OpState, HeavyValueT, Arg);`
`279`	`281`	`}`
`280`	`282`
`281`	`283`	`using namespace mlir;`