[Heavy] Add MemberName as callee to Visit

ricejasonf · ricejasonf · commit eb21b7319fcc · 2025-11-13T14:39:03.000-08:00
diff --git a/heavy/include/nbdl_gen/Nbdl.td b/heavy/include/nbdl_gen/Nbdl.td
@@ -94,21 +94,28 @@ def Nbdl_Variant : TypeDef<Nbdl_Dialect, "Variant", []> {
 */
 }
 
-def Nbdl_Symbol : TypeDef<Nbdl_Dialect, "Symbol"> {
-  let mnemonic = "symbol";
+def Nbdl_MemberName : TypeDef<Nbdl_Dialect, "member_name"> {
+  let mnemonic = "member_name";
   let description = [{
     Represent a string that is a valid C++ identifier.
     The intended primary use case is to specify a member of a Struct.
   }];
 }
 
-// Keys are optional, but we are using variadic arguments
-// so we resort to using this sum type.
-def Nbdl_KeyType : AnyTypeOf<[Nbdl_Store, Nbdl_Tag,
-                              Nbdl_Symbol, Nbdl_Unit]>;
+// TODO Eventually add more specific constraints like variant<...>.
+def Nbdl_Type : AnyTypeOf<[Nbdl_Store,
+                           Nbdl_Unit,
+                           Nbdl_Tag]>;
 
-def Nbdl_Type : AnyTypeOf<[Nbdl_Unit, Nbdl_Empty,
-                           Nbdl_Store, Nbdl_Tag]>;
+// TypePlus is for keys callable objects which allow member names
+// used for accessing members or calling member functions.
+def Nbdl_TypePlus : AnyTypeOf<[Nbdl_Store,
+                               Nbdl_Unit,
+                               Nbdl_Tag,
+                               Nbdl_MemberName]>;
+
+// Implicitly default to the unit type for keys.
+def Nbdl_Key : Optional<Nbdl_TypePlus>;
 
 class Nbdl_Op<string mnemonic, list<Trait> traits = []> :
     Op<Nbdl_Dialect, mnemonic, traits>;
@@ -129,8 +136,14 @@ def Nbdl_GetOp : Nbdl_Op<"get", []> {
   let description = [{
   }];
   let arguments = (ins Nbdl_Type:$state,
-                       Nbdl_KeyType:$key);
+                       Nbdl_Key:$key);
   let results   = (outs Nbdl_Type:$result);
+  let extraClassDeclaration = [{
+    bool hasUnitKey() {
+      return !getKey() ||
+        ::llvm::isa<::nbdl_gen::UnitType>(getKey().getType());
+    }
+  }];
 }
 
 def Nbdl_MatchOp : Nbdl_Op<"match", [Terminator, NoTerminator]> {
@@ -154,9 +167,15 @@ def Nbdl_MatchOp : Nbdl_Op<"match", [Terminator, NoTerminator]> {
   }];
 
   let arguments = (ins Nbdl_Type:$store,
-                       Nbdl_KeyType:$key);
+                       Nbdl_Key:$key);
   let results = (outs);
   let regions = (region SizedRegion<1>:$overloads);
+  let extraClassDeclaration = [{
+    bool hasUnitKey() {
+      return !getKey() ||
+        ::llvm::isa<::nbdl_gen::UnitType>(getKey().getType());
+    }
+  }];
 }
 
 def Nbdl_OverloadOp : Nbdl_Op<"overload", []> {
@@ -204,7 +223,7 @@ def Nbdl_MemberNameOp : Nbdl_Op<"member_name", []> {
   }];
 
   let arguments = (ins StrAttr:$name);
-  let results = (outs Nbdl_Symbol:$result);
+  let results = (outs Nbdl_MemberName:$result);
 }
 
 def Nbdl_StoreComposeOp : Nbdl_Op<"store_compose", []> {
@@ -218,10 +237,16 @@ def Nbdl_StoreComposeOp : Nbdl_Op<"store_compose", []> {
   }];
 
   let arguments = (ins
-    Nbdl_KeyType:$key,
+    Nbdl_Key:$key,
     Nbdl_Type:$lhs,
     Nbdl_Type:$rhs);
   let results = (outs Nbdl_Type:$result);
+  let extraClassDeclaration = [{
+    bool hasUnitKey() {
+      return !getKey() ||
+        ::llvm::isa<::nbdl_gen::UnitType>(getKey().getType());
+    }
+  }];
 }
 
 def Nbdl_VariantOp : Nbdl_Op<"variant", []> {
@@ -353,7 +378,7 @@ def Nbdl_VisitOp : Nbdl_Op<"visit", []> {
     returning a result. This is analogous to std::visit
     except only one function is generated.
   }];
-  let arguments = (ins Nbdl_Type:$fn,
+  let arguments = (ins Nbdl_TypePlus:$fn,
                        Variadic<Nbdl_Type>:$args);
   let results = (outs Nbdl_Type:$result);
 }
diff --git a/heavy/lib/Nbdl/NbdlWriter.cpp b/heavy/lib/Nbdl/NbdlWriter.cpp
@@ -320,7 +320,9 @@ class NbdlWriter {
 
   void VisitType(StoreComposeOp Op) {
     OS << "::nbdl::store_composite<";
-    VisitType(Op.getKey());
+    assert(!Op.hasUnitKey() &&
+      "unit key for store compose currently not supported");
+    VisitTypeOrUnitType(Op.getKey());
     OS << ", ";
     VisitType(Op.getLhs());
     OS << ", ";
@@ -329,14 +331,19 @@ class NbdlWriter {
   }
 
   void VisitType(ConstexprOp Op) {
-    VisitType(Op.getResult());
+    VisitType(Op.getLoc(), Op.getResult().getType());
+  }
+
+  void VisitTypeOrUnitType(mlir::Value V) {
+    assert(V && isa<nbdl_gen::UnitType>(V.getType())
+        && "printing unit type currently not supported");
+
+    if (!V || !isa<nbdl_gen::UnitType>(V.getType()))
+      VisitType(V);
   }
 
   void VisitType(mlir::Location Loc, mlir::Type Type) {
-    if (auto OpaqueType = dyn_cast<nbdl_gen::OpaqueType>(Type))
-      OS << "decltype(auto)";
-    else
-      SetError(Loc, "unsupported type");
+    OS << "auto&&";
   }
 };
 
@@ -412,7 +419,11 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
   }
 
   void Visit(GetOp Op) {
-    auto MemberNameOp = Op.getKey().getDefiningOp<nbdl_gen::MemberNameOp>();
+    bool HasUnitKey = Op.hasUnitKey();
+    nbdl_gen::MemberNameOp MemberNameOp;
+    if (!HasUnitKey)
+      MemberNameOp = Op.getKey().getDefiningOp<nbdl_gen::MemberNameOp>();
+
     OS << "auto&& "
        << SetLocalVarName(Op.getResult(), "get_")
        << " = ";
@@ -423,7 +434,7 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
     } else {
       OS << "::nbdl::get(";
       WriteExpr(Op.getState());
-      if (!isa<nbdl_gen::UnitType>(Op.getKey().getType())) {
+      if (!HasUnitKey) {
         OS << ", ";
         WriteExpr(Op.getKey());
       }
@@ -440,7 +451,20 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
          << " = ";
     }
 
-    WriteExpr(Op.getFn());
+    mlir::Value Fn = Op.getFn();
+    mlir::OperandRange Args = Op.getArgs();
+
+    if (auto MemberNameOp = Fn.getDefiningOp<nbdl_gen::MemberNameOp>()) {
+      if (Args.empty()) {
+        SetError("member literal callee expects at least one argument", Op);
+        return;
+      }
+      WriteExpr(Args.front());
+      OS << '.' << MemberNameOp.getName();
+      Args = Args.drop_front();
+    } else {
+      WriteExpr(Op.getFn());
+    }
     OS << '(';
     llvm::interleave(Op.getArgs(), OS,
         [&](mlir::Value V) {
@@ -458,8 +482,10 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
     OS << "auto ";
     OS << SetLocalVarName(Op.getResult(), "result_");
     OS << " = ::nbdl::store_compose(";
-    WriteExpr(Op.getKey());
-    OS << ",";
+    if (!Op.hasUnitKey()) {
+      WriteExpr(Op.getKey());
+      OS << ",";
+    }
     WriteExpr(Op.getRhs());
     OS << ",";
     WriteExpr(Op.getLhs());
@@ -475,7 +501,7 @@ class FuncWriter : public NbdlWriter<FuncWriter> {
   void Visit(MatchOp Op) {
     OS << "::nbdl::match(";
     WriteExpr(Op.getStore());
-    if (!isa<nbdl_gen::UnitType>(Op.getKey().getType())) {
+    if (!Op.hasUnitKey()) {
       OS << ", ";
       WriteExpr(Op.getKey());
     }
@@ -659,11 +685,12 @@ class DefineStoreWriter : public NbdlWriter<DefineStoreWriter> {
     heavy::SourceLocation PrevLoc = CurLoc;
     SetLoc(Op.getLoc());
 
-    mlir::Value Key = Op.getKey();
+    mlir::Value Key = Op.getKey();  // Could be mlir::Value().
     mlir::Value Lhs = Op.getLhs();
     llvm::StringRef Name;
 
-    if (auto MemberNameOp = Key.getDefiningOp<nbdl_gen::MemberNameOp>()) {
+    if (Key && Key.getDefiningOp<nbdl_gen::MemberNameOp>()) {
+      auto MemberNameOp = Key.getDefiningOp<nbdl_gen::MemberNameOp>();
       // It would be more consistent with our definition of StoreCompose
       // to support shadowing here, but since it is more work to check
       // and very suboptimal for the C++ compiler
diff --git a/heavy/test/Nbdl/close_previous_scope.scm b/heavy/test/Nbdl/close_previous_scope.scm
@@ -9,7 +9,7 @@
 
 (define !nbdl.store (type "!nbdl.store"))
 (define !nbdl.unit (type "!nbdl.unit"))
-(define !nbdl.symbol (type "!nbdl.symbol"))
+(define !nbdl.member_name (type "!nbdl.member_name"))
 (define !nbdl.unit (type "!nbdl.unit"))
 (define i32 (type "i32"))
 
diff --git a/heavy/test/Nbdl/context.scm b/heavy/test/Nbdl/context.scm
@@ -9,7 +9,7 @@
 
 (define !nbdl.store (type "!nbdl.store"))
 (define !nbdl.tag (type "!nbdl.tag"))
-(define !nbdl.symbol (type "!nbdl.symbol"))
+(define !nbdl.member_name (type "!nbdl.member_name"))
 (define !nbdl.unit (type "!nbdl.unit"))
 (define i32 (type "i32"))
 
@@ -20,7 +20,7 @@
       (operands:)
       (attributes:
         ("name" (string-attr name)))
-      (result-types: !nbdl.symbol))))
+      (result-types: !nbdl.member_name))))
 
 (define my_store
   (with-builder
diff --git a/heavy/test/Nbdl/match_params.scm b/heavy/test/Nbdl/match_params.scm
@@ -9,7 +9,7 @@
 
 (define !nbdl.store (type "!nbdl.store"))
 (define !nbdl.tag (type "!nbdl.tag"))
-(define !nbdl.symbol (type "!nbdl.symbol"))
+(define !nbdl.member_name (type "!nbdl.member_name"))
 (define !nbdl.unit (type "!nbdl.unit"))
 (define i32 (type "i32"))
 
@@ -43,7 +43,7 @@
         (old-create-op "nbdl.member_name"
                    (attributes
                      `("name", (string-attr "bar")))
-                   (result-types !nbdl.symbol))))
+                   (result-types !nbdl.member_name))))
     (define key3
       (result
         (old-create-op "nbdl.literal"
