[Heavy] Add %build-context

Author: ricejasonf

heavy/include/nbdl_gen/Nbdl.td

@@ -18,7 +18,6 @@ def Nbdl_Dialect : Dialect {
   let useDefaultTypePrinterParser = 1;
 }
 
-// FIXME Maybe we should just use MLIR's NoneType for this.
 def Nbdl_Empty : TypeDef<Nbdl_Dialect, "Empty", []> {
   let mnemonic = "empty";
   let description = [{
@@ -114,7 +113,7 @@ def Nbdl_Symbol : TypeDef<Nbdl_Dialect, "Symbol"> {
 // Keys are optional, but we are using variadic arguments
 // so we resort to using this sum type.
 def Nbdl_KeyType : AnyTypeOf<[Nbdl_Opaque, Nbdl_Tag,
-                              Nbdl_Symbol, NoneType]>;
+                              Nbdl_Symbol, Nbdl_Empty]>;
 
 def Nbdl_Type : AnyTypeOf<[Nbdl_Void, Nbdl_Empty, Nbdl_Opaque,
                            Nbdl_Store, Nbdl_Tag]>;
@@ -236,7 +235,7 @@ def Nbdl_StoreComposeOp : Nbdl_Op<"store_compose", []> {
   }];
 
   let arguments = (ins
-    TypeAttr:$key,
+    Nbdl_KeyType:$key,
     Nbdl_Type:$lhs,
     Nbdl_Type:$rhs);
   let results = (outs Nbdl_Type:$result);
@@ -250,7 +249,7 @@ def Nbdl_VariantOp : Nbdl_Op<"variant", []> {
   let results = (outs Nbdl_Type:$result);
 }
 
-def Nbdl_CreateStoreOp : Nbdl_Op<"create_store", [Symbol, IsolatedFromAbove]> {
+def Nbdl_ContextOp : Nbdl_Op<"context", [Symbol, IsolatedFromAbove]> {
   let description = [{
     Define a store object, and expose its interface (type) to the user.
 
@@ -270,9 +269,16 @@ def Nbdl_CreateStoreOp : Nbdl_Op<"create_store", [Symbol, IsolatedFromAbove]> {
   }];
 
   let regions = (region AnyRegion:$body);
-  let arguments = (ins
-    FlatSymbolRefAttr:$name,
-    OptionalAttr<DictArrayAttr>:$cnstrArgs);
+  let arguments = (ins StrAttr:$sym_name);
+  let results = (outs);
+}
+
+def Nbdl_ContOp : Nbdl_Op<"cont", [Terminator]> {
+  let description = [{
+    The terminator we need.
+  }];
+
+  let arguments = (ins Nbdl_Type:$arg);
   let results = (outs);
 }
 
@@ -289,6 +295,15 @@ def Nbdl_StoreOp : Nbdl_Op<"store", []> {
   let results = (outs Nbdl_Type:$result);
 }
 
+def Nbdl_EmptyOp : Nbdl_Op<"empty", []> {
+  let description = [{
+    An operation that produces an object of the empty type.
+  }];
+
+  let arguments = (ins);
+  let results = (outs Nbdl_Empty:$result);
+}
+
 def Nbdl_ApplyOp : Nbdl_Op<"apply", []> {
   let description = [{
     Apply a function.

heavy/lib/Nbdl.cpp

@@ -32,6 +32,7 @@ heavy::ExternFunction translate_cpp;
 heavy::ExternFunction build_match_params_impl;
 heavy::ExternFunction build_overload_impl;
 heavy::ExternFunction build_match_if_impl;
+heavy::ExternFunction build_context_impl;
 }
 
 namespace {
@@ -226,7 +227,63 @@ void translate_cpp(Context& C, ValueRefs Args) {
   }
   C.Cont();
 }
+
+void build_context_impl(Context& C, ValueRefs Args) {
+  if (Args.size() != 3)
+    return C.RaiseError("invalid arity");
+
+  std::optional<mlir::OpBuilder> BuilderOpt = getModuleBuilder(C);
+  if (!BuilderOpt)
+    return;
+  mlir::OpBuilder Builder = BuilderOpt.value();
+
+  // This part is exactly like build_match_params_impl
+  heavy::SourceLocation Loc = Args[0].getSourceLocation();
+  // Require a heavy::Symbol so it has a source location.
+  heavy::Symbol* NameSym = dyn_cast<heavy::Symbol>(Args[0]);
+  llvm::StringRef Name = NameSym->getStringRef();
+  heavy::Value NumParamsVal = Args[1];
+  heavy::Value Callback = Args[2];
+
+  int32_t NumParams = isa<heavy::Int>(NumParamsVal)
+      ? int32_t(cast<heavy::Int>(NumParamsVal))
+      : int32_t(-1);
+
+  if (NumParams < 0)
+    return C.RaiseError("expecting positive integer for num_params");
+  if (!NameSym || Name.empty())
+    return C.RaiseError("expecting function name (symbol literal)");
+
+  mlir::Location MLoc = mlir::OpaqueLoc::get(Loc.getOpaqueEncoding(),
+                                             Builder.getContext());
+  // Create the ContextOp.
+  auto ContextOp = Builder.create<nbdl_gen::ContextOp>(MLoc, Name);
+  mlir::Block& EntryBlock = ContextOp.getBody().emplaceBlock();
+
+  // Create the arguments.
+  auto OpaqueTy = Builder.getType<nbdl_gen::OpaqueType>();
+  for (int32_t i = 0; i < NumParams; i++)
+    EntryBlock.addArgument(OpaqueTy, MLoc);
+
+  heavy::Value Thunk = C.CreateLambda([ContextOp](Context& C,
+                                                  ValueRefs) mutable {
+    heavy::Value Callback = C.getCapture(0);
+    llvm::SmallVector<heavy::Value, 8> BlockArgs;
+    assert(!ContextOp.getBody().empty() && "should have entry block");
+    for (mlir::Value MVal : ContextOp.getBody().getArguments()) {
+      heavy::Value V = mlir_helper::createTagged(C,
+          mlir_helper::kind::mlir_value, MVal);
+      BlockArgs.push_back(V);
+    }
+
+    C.Apply(Callback, BlockArgs);
+  }, CaptureList{Callback});
+
+  // Call the thunk with a Builder at the entry point.
+  Builder = mlir::OpBuilder(ContextOp.getBody());
+  mlir_helper::with_builder_impl(C, Builder, Thunk);
 }
+} //  end namespace heavy::nbdl_bind
 
 extern "C" {
 // initialize the module for run-time independent of the compiler
@@ -247,6 +304,8 @@ void HEAVY_NBDL_INIT(heavy::Context& C) {
     = heavy::nbdl_bind::build_overload_impl;
   heavy::nbdl_bind_var::build_match_if_impl
     = heavy::nbdl_bind::build_match_if_impl;
+  heavy::nbdl_bind_var::build_context_impl
+    = heavy::nbdl_bind::build_context_impl;
 }
 
 void HEAVY_NBDL_LOAD_MODULE(heavy::Context& C) {
@@ -257,6 +316,7 @@ void HEAVY_NBDL_LOAD_MODULE(heavy::Context& C) {
     {"%build-match-params", heavy::nbdl_bind_var::build_match_params_impl},
     {"%build-overload", heavy::nbdl_bind_var::build_overload_impl},
     {"%build-match-if", heavy::nbdl_bind_var::build_match_if_impl},
+    {"%build-context", heavy::nbdl_bind_var::build_context_impl},
   });
 }
 }

heavy/lib/NbdlWriter.cpp

@@ -45,7 +45,7 @@ class NbdlWriter {
   heavy::SourceLocationEncoding* ErrLoc;
   llvm::raw_ostream& OS;
 
-  // Track number of members for CreateStoreOp
+  // Track number of members of context
   // to generate anonymous identifiers if needed.
   unsigned CurrentMemberCount = 0;
   unsigned CurrentAnonVarCount = 0;
@@ -134,7 +134,7 @@ class NbdlWriter {
     if (CheckError())
       return;
 
-         if (isa<CreateStoreOp>(Op))  return Visit(cast<CreateStoreOp>(Op));
+         if (isa<ContextOp>(Op))      return Visit(cast<ContextOp>(Op));
     else if (isa<StoreOp>(Op))        return Visit(cast<StoreOp>(Op));
     else if (isa<ApplyOp>(Op))        return Visit(cast<ApplyOp>(Op));
     else if (isa<GetOp>(Op))          return Visit(cast<GetOp>(Op));
@@ -161,7 +161,7 @@ class NbdlWriter {
       Visit(&Op);
   }
 
-  void Visit(CreateStoreOp Op) {
+  void Visit(ContextOp Op) {
     // Skip externally defined stores.
     if (Op.isExternal())
       return;
@@ -197,7 +197,7 @@ class NbdlWriter {
 #if 0
     if (isTopLevel()) {
       // Add the RHS as a member.
-      assert(isa<nbdl_gen::CreateStoreOp>(Op.getParentOp()) &&
+      assert(isa<nbdl_gen::ContextOp>(Op.getParentOp()) &&
              "should be in context of creating a store");
 
       // Temporarily store anonymous member name if needed.
@@ -385,8 +385,8 @@ class NbdlWriter {
    ************************************/
 
   void VisitType(mlir::Operation* Op) {
-         if (isa<CreateStoreOp>(Op))
-           return VisitType(cast<CreateStoreOp>(Op));
+         if (isa<ContextOp>(Op))
+           return VisitType(cast<ContextOp>(Op));
     else if (isa<StoreOp>(Op))    return VisitType(cast<StoreOp>(Op));
     else if (isa<VariantOp>(Op))      return VisitType(cast<VariantOp>(Op));
     else if (isa<StoreComposeOp>(Op))
@@ -396,7 +396,7 @@ class NbdlWriter {
       SetError("unhandled operation (VisitType)", Op);
   }
 
-  void VisitType(CreateStoreOp Op) {
+  void VisitType(ContextOp Op) {
     OS << Op.getName();
   }
 
@@ -415,6 +415,8 @@ class NbdlWriter {
   }
 
   void VisitType(StoreComposeOp Op) {
+    llvm_unreachable("TODO");
+#if 0
     OS << "nbdl::store_composite<";
 
     // Key
@@ -427,6 +429,7 @@ class NbdlWriter {
     OS << ", ";
     VisitType(Op.getLoc(), Op.getLhs());
     OS << ">";
+#endif
   }
 
   void VisitType(ConstexprOp Op) {