[Heavy] Support multivariate continuations. Fix tail calls.

ricejasonf · ricejasonf · commit 36fe13bed147 · 2025-09-05T21:58:25.000-07:00
diff --git a/heavy/include/heavy/ContinuationStack.h b/heavy/include/heavy/ContinuationStack.h
@@ -135,7 +135,7 @@ class ContinuationStack {
     DidCallContinuation = true; // debug mode only
     if (isa<Undefined>(Callee))
       return getDerived().RaiseError("callee is undefined");
-    if (Args.data() != ApplyArgs.data()) {
+    if (Args.data() != ValueRefs(ApplyArgs).drop_front().data()) {
       ApplyArgs.resize(Args.size() + 1);
       std::copy(Args.begin(), Args.end(), ApplyArgs.begin() + 1);
     }
@@ -194,7 +194,10 @@ class ContinuationStack {
     assert(ApplyArgs[0] && "callee must not be null");
     return ApplyArgs[0];
   }
-  ValueRefs getCaptures() {
+  heavy::ValueRefs getCallArgs() {
+    return heavy::ValueRefs(ApplyArgs).drop_front();
+  }
+  heavy::ValueRefs getCaptures() {
     return cast<Lambda>(ApplyArgs[0])->getCaptures();
   }
   heavy::Value getCapture(unsigned I) {
@@ -243,7 +246,7 @@ class ContinuationStack {
       // Debug mode only
       DidCallContinuation = false;
 
-      ValueRefs Args = ValueRefs(ApplyArgs).drop_front();
+      ValueRefs Args = getCallArgs();
       switch (Callee.getKind()) {
       case ValueKind::Lambda: {
         Lambda* L = cast<Lambda>(Callee);
@@ -302,10 +305,6 @@ class ContinuationStack {
   void PushCont(heavy::Value Callable) {
     PushCont([](Derived& Context, ValueRefs Args) {
       heavy::Value Callable = Context.getCapture(0);
-      // FIXME Remove this drop_front when is certain
-      //       that it was vestigal
-      //       (we used to include the callee in Args)
-      //Context.Apply(Callable, Args.drop_front());
       Context.Apply(Callable, Args);
     }, Callable);
   }
@@ -381,21 +380,21 @@ class ContinuationStack {
     C.Apply(InputProc, Proc);
   }
 
-  // SaveEscapeProc - Push Proc as current continuation for use 
+  // SaveEscapeProc - Push Proc as current continuation for use
   //                  as an escape procedure bound to Var, then
   //                  call the continuation we had before so
   //                  that Proc is only called explicitly.
   template <typename F>
   void SaveEscapeProc(Value Var, F Proc, CaptureList Captures) {
-    assert(isa<Binding>(Var) && "expecting a binding for Var"); 
+    assert(isa<Binding>(Var) && "expecting a binding for Var");
     Derived& C = getDerived();
     PushCont(Proc, Captures);
     CallCC(C.CreateLambda([](Derived& C, ValueRefs Args) {
       cast<Binding>(C.getCapture(0))->setValue(Args[0]);
       // Remove Proc from the stack before continuing.
       C.PopCont();
       C.Cont();
-    }, CaptureList{Var})); 
+    }, CaptureList{Var}));
   }
 
   void DynamicWind(Value Before, Value Thunk, Value After) {
@@ -422,7 +421,7 @@ class ContinuationStack {
         C.PushCont([](Derived& C, ValueRefs) {
           ValueRefs ThunkResults = C.getCaptures();
           C.Cont(ThunkResults);
-        }, /*Captures=*/ThunkResults); 
+        }, /*Captures=*/ThunkResults);
         C.CurDW = PrevDW;
         C.Apply(After, {});
       }, CaptureList{After, PrevDW});
diff --git a/heavy/include/heavy/Dialect.h b/heavy/include/heavy/Dialect.h
@@ -51,6 +51,15 @@ struct Dialect : public mlir::Dialect {
   void printType(mlir::Type, mlir::DialectAsmPrinter&) const override;
 };
 
+struct HeavyContextTy : public mlir::Type::TypeBase<
+                            HeavyContextTy,
+                            mlir::Type,
+                            mlir::TypeStorage> {
+  static constexpr llvm::StringLiteral name = "heavy.context";
+  static constexpr llvm::StringRef getMnemonic() { return "context"; }
+  using Base::Base;
+};
+
 struct HeavyValueTy : public mlir::Type::TypeBase<
                             HeavyValueTy,
                             mlir::Type,
@@ -60,6 +69,19 @@ struct HeavyValueTy : public mlir::Type::TypeBase<
   using Base::Base;
 };
 
+// Represent a variadiac array of arguments for use with
+// continuation arguments.
+struct HeavyValueRefsTy : public mlir::Type::TypeBase<
+                            HeavyValueRefsTy,
+                            mlir::Type,
+                            mlir::TypeStorage> {
+  static constexpr llvm::StringLiteral name = "heavy.value_refs";
+  static constexpr llvm::StringRef getMnemonic() { return "value_refs"; }
+  using Base::Base;
+};
+
+// Represent a variadic list of arguments which are provided
+// as a scheme linked list.
 struct HeavyRestTy : public mlir::Type::TypeBase<
                             HeavyRestTy,
                             mlir::Type,
diff --git a/heavy/include/heavy/OpGen.h b/heavy/include/heavy/OpGen.h
@@ -257,6 +257,7 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
       return heavy::Value(OpResult.getOwner());
     }
     if (auto BlockArg = mlir::dyn_cast<mlir::BlockArgument>(V)) {
+      // FIXME Do we even get here? Is ContArg ever needed?
       mlir::Block* B = BlockArg.getOwner();
       return heavy::Value(reinterpret_cast<heavy::ContArg*>(B));
     }
@@ -276,8 +277,14 @@ class OpGen : public ValueVisitor<OpGen, mlir::Value> {
                             Value Else);
   mlir::Value createContinuation(mlir::Operation* CallOp);
 
+  enum class RestParamKind {
+    None = 0,
+    List,
+    ValueRefs,
+  };
+
   mlir::FunctionType createFunctionType(unsigned Arity,
-                                        bool HasRestParam);
+                                        RestParamKind RPK);
   heavy::FuncOp createFunction(SourceLocation Loc,
                                llvm::StringRef MangledName,
                                mlir::FunctionType FT);
diff --git a/heavy/include/heavy/Ops.td b/heavy/include/heavy/Ops.td
@@ -40,7 +40,9 @@ def HeavyValueAttr : AttrDef<HeavyDialect, "HeavyValue"> {
 }
 
 // Additional Types
-def HeavyRest : HeavyValueBase<"Pair">; // For rest arguments.
+def HeavyContext : HeavyValueBase<"Context">;
+def HeavyValueRefs : HeavyValueBase<"ValueRefs">; // For arguments array.
+def HeavyRest : HeavyValueBase<"Rest">; // For rest arguments.
 def HeavyPair : HeavyValueBase<"Pair">;
 def HeavySyntax : HeavyValueBase<"Syntax">;
 
@@ -59,7 +61,8 @@ def heavy_ApplyOp : HeavyOp<"apply", [ReturnLike, Terminator]> {
     a continuation.
   }];
 
-  let arguments = (ins HeavyValue:$fn, Variadic<HeavyValue>:$args);
+  let arguments = (ins HeavyValue:$fn,
+    Variadic<AnyTypeOf<[HeavyValue, HeavyValueRefs]>>:$args);
   let results   = (outs);
 }
 
@@ -266,13 +269,16 @@ def heavy_LiteralOp : HeavyOp<"literal", [Pure]> {
   ];
 }
 
-def heavy_LoadClosureOp : HeavyOp<"load_closure"> {
-  let summary = "load_closure";
+def heavy_LoadRefOp : HeavyOp<"load_ref"> {
+  let summary = "load_ref";
   let description = [{
-    Read an element from lambda's capture list.
+    Read an element from a ValueRefs object which can be
+    the context's callee capture list or a variadic argument
+    list which is primarily used for continuations.
   }];
 
-  let arguments = (ins HeavyValue:$closure, UI32Attr:$index);
+  let arguments = (ins AnyTypeOf<[HeavyContext, HeavyValueRefs]>:$value_refs,
+                       UI32Attr:$index);
   let results   = (outs HeavyValue:$result);
 
   let builders = [
diff --git a/heavy/lib/Builtins.cpp b/heavy/lib/Builtins.cpp
@@ -67,6 +67,8 @@ heavy::ExternFunction eq;
 heavy::ExternFunction equal;
 heavy::ExternFunction eqv;
 heavy::ExternFunction call_cc;
+heavy::ExternFunction values;
+heavy::ExternFunction call_with_values;
 heavy::ExternFunction with_exception_handler;
 heavy::ExternFunction raise;
 heavy::ExternFunction error;
@@ -385,6 +387,19 @@ void call_cc(Context& C, ValueRefs Args) {
   C.CallCC(Args[0]);
 }
 
+void values(Context& C, ValueRefs Args) {
+  C.Cont(Args);
+}
+
+void call_with_values(Context& C, ValueRefs Args) {
+  if (Args.size() != 2)
+    return C.RaiseError("invalid arity");
+  heavy::Value Producer = Args[0];
+  heavy::Value Consumer = Args[1];
+  C.PushCont(Consumer);
+  C.Apply(Producer, {});
+}
+
 void dump(Context& C, ValueRefs Args) {
   if (Args.size() != 1) return C.RaiseError("invalid arity");
   Args[0].dump();
@@ -888,6 +903,8 @@ void HEAVY_BASE_INIT(heavy::Context& Context) {
   HEAVY_BASE_VAR(equal)   = heavy::base::equal;
   HEAVY_BASE_VAR(eqv)     = heavy::base::eqv;
   HEAVY_BASE_VAR(call_cc) = heavy::base::call_cc;
+  HEAVY_BASE_VAR(values)  = heavy::base::values;
+  HEAVY_BASE_VAR(call_with_values) = heavy::base::call_with_values;
   HEAVY_BASE_VAR(with_exception_handler)
     = heavy::base::with_exception_handler;
   HEAVY_BASE_VAR(raise)   = heavy::base::raise;
@@ -965,6 +982,8 @@ void HEAVY_BASE_LOAD_MODULE(heavy::Context& Context) {
     {"equal?",  HEAVY_BASE_VAR(equal)},
     {"eqv?",    HEAVY_BASE_VAR(eqv)},
     {"call/cc", HEAVY_BASE_VAR(call_cc)},
+    {"values", HEAVY_BASE_VAR(values)},
+    {"call-with-values", HEAVY_BASE_VAR(call_with_values)},
     {"with-exception-handler", HEAVY_BASE_VAR(with_exception_handler)},
     {"raise", HEAVY_BASE_VAR(raise)},
     {"error", HEAVY_BASE_VAR(error)},
diff --git a/heavy/lib/Dialect.cpp b/heavy/lib/Dialect.cpp
@@ -27,6 +27,8 @@ Dialect::Dialect(mlir::MLIRContext* Ctx)
   addTypes<HeavyValueTy>();
   addAttributes<HeavyValueAttr>();
 
+  addTypes<HeavyContextTy>();
+  addTypes<HeavyValueRefsTy>();
   addTypes<HeavyRestTy>();
   addTypes<HeavyPairTy>();
 
@@ -58,6 +60,12 @@ mlir::Type Dialect::parseType(mlir::DialectAsmParser& P) const {
   mlir::Builder B = P.getBuilder();
   if (Name == HeavyValueTy::getMnemonic())
     return B.getType<HeavyValueTy>();
+  if (Name == HeavyContextTy::getMnemonic())
+    return B.getType<HeavyContextTy>();
+  if (Name == HeavyValueRefsTy::getMnemonic())
+    return B.getType<HeavyValueRefsTy>();
+  if (Name == HeavyRestTy::getMnemonic())
+    return B.getType<HeavyRestTy>();
   if (Name == HeavySyntaxTy::getMnemonic())
     return B.getType<HeavySyntaxTy>();
   if (Name == HeavyPairTy::getMnemonic())
@@ -81,6 +89,10 @@ void Dialect::printType(mlir::Type Type,
   char const* Name;
   if (mlir::isa<HeavyValueTy>(Type)) {
     Name = "value";
+  } else if (mlir::isa<HeavyContextTy>(Type)) {
+    Name = "context";
+  } else if (mlir::isa<HeavyValueRefsTy>(Type)) {
+    Name = "value_refs";
   } else if (mlir::isa<HeavyRestTy>(Type)) {
     Name = "rest";
   } else if (mlir::isa<HeavyPairTy>(Type)) {
@@ -147,10 +159,10 @@ void LiteralOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
                    HeavyValueAttr::get(B.getContext(), V));
 }
 
-void LoadClosureOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
-                          mlir::Value Closure, uint32_t Index) {
+void LoadRefOp::build(mlir::OpBuilder& B, mlir::OperationState& OpState,
+                          mlir::Value ValueRefs, uint32_t Index) {
   mlir::Type HeavyValueT = B.getType<HeavyValueTy>();
-  LoadClosureOp::build(B, OpState, HeavyValueT, Closure,
+  LoadRefOp::build(B, OpState, HeavyValueT, ValueRefs,
                        B.getUI32IntegerAttr(Index));
 }
 
diff --git a/heavy/lib/OpEval.cpp b/heavy/lib/OpEval.cpp
diff --git a/heavy/lib/OpGen.cpp b/heavy/lib/OpGen.cpp
diff --git a/heavy/test/Evaluate/values.scm b/heavy/test/Evaluate/values.scm
