[Heavy] Add DefineStoreOp

Author: ricejasonf

heavy/include/heavy/Value.h

@@ -1399,10 +1399,17 @@ class Any final :
 
   // Note that for stored pointers, this will return a pointer to a pointer.
   void* getOpaquePtr() { return getTrailingObjects<char>(); }
+
   llvm::StringRef getObjData() {
     return llvm::StringRef(static_cast<char*>(getOpaquePtr()), getObjectSize());
   };
 
+  bool equal(Any* Other) {
+    return TypeId == Other->TypeId &&
+           StorageLen == Other->StorageLen &&
+           getObjData() == Other->getObjData();
+  }
+
   template <typename T>
   bool isa() {
     return TypeId == &AnyTypeId<T>::Id;

heavy/include/nbdl_gen/Nbdl.td

@@ -71,13 +71,6 @@ def Nbdl_Store : TypeDef<Nbdl_Dialect, "Store", []> {
   }];
 }
 
-def Nbdl_Variant : TypeDef<Nbdl_Dialect, "Variant", []> {
-  let mnemonic = "variant";
-  let description = [{
-    A type satisfying the nbdl::Variant concept
-  }];
-}
-
 def Nbdl_Tag : TypeDef<Nbdl_Dialect, "Tag", []> {
   let mnemonic = "tag";
   let description = [{
@@ -86,6 +79,21 @@ def Nbdl_Tag : TypeDef<Nbdl_Dialect, "Tag", []> {
   }];
 }
 
+def Nbdl_Variant : TypeDef<Nbdl_Dialect, "Variant", []> {
+  let mnemonic = "variant";
+  let description = [{
+    A type satisfying the nbdl::Variant concept
+  }];
+
+/*
+  // TODO Eventually it would be nice to specify the alternatives
+  //      Also some parameterization for knowing the key => 
+  let parameters = (ins Variadic<AnyTypeOf<[Nbdl_Store, Nbdl_Tag,
+                                            Nbdl_Unit>:$alts);
+  let assemblyFormat = "`<`$alts`>`";
+*/
+}
+
 def Nbdl_Symbol : TypeDef<Nbdl_Dialect, "Symbol"> {
   let mnemonic = "symbol";
   let description = [{
@@ -236,7 +244,7 @@ def Nbdl_VariantOp : Nbdl_Op<"variant", []> {
   let results = (outs Nbdl_Type:$result);
 }
 
-def Nbdl_ContextOp : Nbdl_Op<"context", [Symbol, IsolatedFromAbove]> {
+def Nbdl_DefineStoreOp : Nbdl_Op<"define_store", [Symbol, IsolatedFromAbove]> {
   let description = [{
     Define a store object, and expose its interface (type) to the user.
 
@@ -268,6 +276,33 @@ def Nbdl_ContextOp : Nbdl_Op<"context", [Symbol, IsolatedFromAbove]> {
   }];
 }
 
+def Nbdl_ContextOp : Nbdl_Op<"context", [Symbol]> {
+  let description = [{
+    Encapsulate a store object to serve as a root node in a state tree.
+
+    Nested objects are still accessible via `match` and `get`,
+    but the context wrapper supports the convention that the root
+    node cannot be a non-root member of a path.
+
+    An instance can be created with nbdl.store.
+  }];
+
+  let arguments = (ins StrAttr:$sym_name, FlatSymbolRefAttr:$implName);
+  let results = (outs);
+
+  let extraClassDeclaration = [{
+    // Get the fully qualified name.
+    // This is needed for defining template specializations.
+    llvm::StringRef getFullName() { return getSymName(); }
+
+    // Get the unqualified name.
+    llvm::StringRef getName() {
+      llvm::StringRef SymName = getSymName();
+      return SymName.take_back(SymName.rfind(':'));
+    }
+  }];
+}
+
 def Nbdl_ContOp : Nbdl_Op<"cont", [Terminator]> {
   let description = [{
     The terminator we need.

heavy/lib/Context.cpp

@@ -727,6 +727,8 @@ bool eqv_slow(Value V1, Value V2) {
   case ValueKind::Float:
     return cast<Float>(V1)->getVal() ==
               cast<Float>(V2)->getVal();
+  case ValueKind::Any:
+    return cast<Any>(V1)->equal(cast<Any>(V2));
   default:
       return false;
   }

heavy/lib/Nbdl.cpp

@@ -215,7 +215,7 @@ void translate_cpp(Context& C, ValueRefs Args) {
     return C.RaiseError("invalid arity");
   auto* Op = dyn_cast<mlir::Operation>(Args[0]);
   if (!Op)
-    return C.RaiseError("expecting mlir.operation");
+    return C.RaiseError("expecting mlir.operation: {}", Args[0]);
 
   llvm::raw_ostream* OS = nullptr;
 
@@ -285,6 +285,8 @@ void close_previous_scope(Context& C, ValueRefs Args) {
 }
 
 void build_context_impl(Context& C, ValueRefs Args) {
+  return C.RaiseError("deprecated");
+#if 0
   if (Args.size() != 3)
     return C.RaiseError("invalid arity");
 
@@ -337,6 +339,7 @@ void build_context_impl(Context& C, ValueRefs Args) {
   // Call the thunk with a Builder at the entry point.
   Builder = mlir::OpBuilder(ContextOp.getBody());
   mlir_helper::with_builder_impl(C, Builder, Thunk);
+#endif
 }
 
 

heavy/lib/Nbdl/NbdlWriter.cpp

@@ -301,7 +301,7 @@ class NbdlWriter {
   }
 
   void VisitType(ContextOp Op) {
-    OS << Op.getName();
+    OS << Op.getSymName();
   }
 
   void VisitType(StoreOp Op) {
@@ -560,43 +560,92 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
 
 class ContextWriter : public NbdlWriter<ContextWriter> {
 public:
-  // Record the values for each member in order.
-  llvm::SmallVector<mlir::Value, 8> Members;
-
   using NbdlWriter<ContextWriter>::NbdlWriter;
 
   void VisitContext(ContextOp Op) {
     SetLoc(Op.getLoc());
+
+    llvm::StringRef ClassName = Op.getName();
+    llvm::StringRef ImplName = Op.getImplName();
+
+    OS << "class " << ClassName << " : nbdl::context_alias<";
+    OS << ImplName;
+    OS << ", /*is_moveable=*/false> {\n"
+          "using Base = ";
+    OS << ImplName;
+    OS << ";\n using Base::Base;\n";
+    OS << "};\n";
+    Flush();
+  }
+};
+
+class DefineStoreWriter : public NbdlWriter<DefineStoreWriter> {
+public:
+  // Record the values for each member in order.
+  llvm::SmallVector<mlir::Value, 8> Members;
+
+  using NbdlWriter<DefineStoreWriter>::NbdlWriter;
+
+  void VisitDefineStore(DefineStoreOp Op) {
     // Skip externally defined stores.
     if (Op.isExternal())
       return;
 
+    SetLoc(Op.getLoc());
     ValueMapScope Scope(ValueMap);
 
     // Set the arg names first.
     for (mlir::BlockArgument BlockArg : Op.getBody().getArguments())
       SetLocalVarName(BlockArg, "arg_");
 
-    // Delete both copy constructors to support subsumption with `auto&&`.
-    OS << "class " << Op.getName() << " {\n";
+    auto ContOp = getContOp(Op);
+    llvm::TypeSwitch<mlir::Operation*>(ContOp.getArg().getDefiningOp())
+      .Case<UnitOp>([&, this](auto) {
+          this->CreateTag(Op.getName());
+        })
+      .Case<StoreOp, VariantOp>([&, this](auto ResultOp) {
+          this->CreateStrongAlias(Op.getName(), ResultOp.getResult());
+        })
+      .Case<StoreComposeOp>([&, this](auto) {
+          this->CreateClass(Op);
+        });
+  }
+
+  void CreateTag(llvm::StringRef Name) {
+    OS << "struct " << Name << " { };\n";
+  }
+
+  void CreateStrongAlias(llvm::StringRef Name, mlir::Value V) {
+    OS << "class " << Name << " : public nbdl::strong_alias<";
+    VisitType(V);
+    OS << "> {\n"
+          "using Base = ";
+    VisitType(V);
+    OS << ";\n using Base::Base;\n";
+    OS << "};\n";
+  }
+
+  void CreateClass(DefineStoreOp Op) {
+    llvm::StringRef Name = Op.getName();
+    OS << "class " << Name << " {\n";
     OS << "public:\n";
     WriteMemberDecls(Op);
-    OS << Op.getName() << '(' << Op.getName() << " const&) = delete;\n";
-    OS << Op.getName() << '(' << Op.getName() << "&) = delete;\n";
+    OS << Name << "(" << Name << " const&) = default;\n";
+    OS << Name << "(" << Name << "&&) = default;\n";
     WriteConstructor(Op);
     OS << "};\n";
     Flush();
   }
 
-  nbdl_gen::ContOp getContOp(ContextOp Op) {
+  nbdl_gen::ContOp getContOp(DefineStoreOp Op) {
     mlir::Operation* Terminator = Op.getBody().front().getTerminator();
     auto ContOp = dyn_cast<nbdl_gen::ContOp>(Terminator);
     if (!ContOp)
       SetError("expecting nbdl.cont as terminator", Op);
     return ContOp;
   }
 
-  void WriteMemberDecls(ContextOp Op) {
+  void WriteMemberDecls(DefineStoreOp Op) {
     // Get the ContOp and work backwards
     // saving the member names as we go.
     auto ContOp = getContOp(Op);
@@ -651,12 +700,12 @@ class ContextWriter : public NbdlWriter<ContextWriter> {
     CurLoc = PrevLoc;
   }
 
-  void WriteConstructor(ContextOp Op) {
+  void WriteConstructor(DefineStoreOp Op) {
     auto ContOp = getContOp(Op);
     if (!ContOp)
       return;
 
-    OS << Op.getName();
+    OS << "explicit " << Op.getName();
     OS << '(';
     llvm::interleaveComma(Op.getBody().getArguments(), OS,
         [&](mlir::BlockArgument const& Arg) {
@@ -748,6 +797,11 @@ translate_cpp(heavy::LexerWriterFnRef LexerWriter, mlir::Operation* Op) {
     Writer.Visit(Op);
     return std::make_tuple(std::move(Writer.ErrMsg),
                            Writer.ErrLoc, Writer.Irritant);
+  } else if (auto DefineStoreOp = dyn_cast<nbdl_gen::DefineStoreOp>(Op)) {
+    DefineStoreWriter Writer(LexerWriter);
+    Writer.VisitDefineStore(DefineStoreOp);
+    return std::make_tuple(std::move(Writer.ErrMsg),
+                           Writer.ErrLoc, Writer.Irritant);
   } else if (auto ContextOp = dyn_cast<nbdl_gen::ContextOp>(Op)) {
     ContextWriter Writer(LexerWriter);
     Writer.VisitContext(ContextOp);

heavy/test/Nbdl/apply_action.scm

@@ -39,6 +39,7 @@
             (operands: Store Num42)
             (attributes:)
             (result-types: !nbdl.store))))
+      (close-previous-scope) ; Should do nothing.
       (create-op "nbdl.match"
         (loc: 0)
         (operands: Store UnitKey)
@@ -59,6 +60,7 @@
                 (result-types:))))))
       ; Foo is not allowed after this (because it could be invalidated.)
       (close-previous-scope)
+      (close-previous-scope) ; Should do nothing.
       (let ((SomeTag (result
                         (create-op "nbdl.constexpr"
                           (loc: 0)