1 files changed, 89 insertions, 23 deletions

diff --git a/clang/lib/Sema/SemaHLSL.cpp b/clang/lib/Sema/SemaHLSL.cpp
index 09e5d69..17cb1e4 100644
--- a/clang/lib/Sema/SemaHLSL.cpp
+++ b/clang/lib/Sema/SemaHLSL.cpp
@@ -1240,6 +1240,20 @@ static CXXMethodDecl *lookupMethod(Sema &S, CXXRecordDecl *RecordDecl,
 
 } // end anonymous namespace
 
+static bool hasCounterHandle(const CXXRecordDecl *RD) {
+  if (RD->field_empty())
+    return false;
+  auto It = std::next(RD->field_begin());
+  if (It == RD->field_end())
+    return false;
+  const FieldDecl *SecondField = *It;
+  if (const auto *ResTy =
+          SecondField->getType()->getAs<HLSLAttributedResourceType>()) {
+    return ResTy->getAttrs().IsCounter;
+  }
+  return false;
+}
+
 bool SemaHLSL::handleRootSignatureElements(
     ArrayRef<hlsl::RootSignatureElement> Elements) {
   // Define some common error handling functions
@@ -2973,6 +2987,25 @@ bool SemaHLSL::CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall) {
     TheCall->setType(ResourceTy);
     break;
   }
+  case Builtin::BI__builtin_hlsl_resource_counterhandlefromimplicitbinding: {
+    ASTContext &AST = SemaRef.getASTContext();
+    if (SemaRef.checkArgCount(TheCall, 3) ||
+        CheckResourceHandle(&SemaRef, TheCall, 0) ||
+        CheckArgTypeMatches(&SemaRef, TheCall->getArg(1), AST.UnsignedIntTy) ||
+        CheckArgTypeMatches(&SemaRef, TheCall->getArg(2), AST.UnsignedIntTy))
+      return true;
+
+    QualType MainHandleTy = TheCall->getArg(0)->getType();
+    auto *MainResType = MainHandleTy->getAs<HLSLAttributedResourceType>();
+    auto MainAttrs = MainResType->getAttrs();
+    assert(!MainAttrs.IsCounter && "cannot create a counter from a counter");
+    MainAttrs.IsCounter = true;
+    QualType CounterHandleTy = AST.getHLSLAttributedResourceType(
+        MainResType->getWrappedType(), MainResType->getContainedType(),
+        MainAttrs);
+    TheCall->setType(CounterHandleTy);
+    break;
+  }
   case Builtin::BI__builtin_hlsl_and:
   case Builtin::BI__builtin_hlsl_or: {
     if (SemaRef.checkArgCount(TheCall, 2))
@@ -3780,10 +3813,24 @@ void SemaHLSL::ActOnVariableDeclarator(VarDecl *VD) {
         uint32_t OrderID = getNextImplicitBindingOrderID();
         if (Binding.hasBinding())
           Binding.setImplicitOrderID(OrderID);
-        else
+        else {
           addImplicitBindingAttrToDecl(
               SemaRef, VD, getRegisterType(getResourceArrayHandleType(VD)),
               OrderID);
+          // Re-create the binding object to pick up the new attribute.
+          Binding = ResourceBindingAttrs(VD);
+        }
+      }
+
+      // Get to the base type of a potentially multi-dimensional array.
+      QualType Ty = getASTContext().getBaseElementType(VD->getType());
+
+      const CXXRecordDecl *RD = Ty->getAsCXXRecordDecl();
+      if (hasCounterHandle(RD)) {
+        if (!Binding.hasCounterImplicitOrderID()) {
+          uint32_t OrderID = getNextImplicitBindingOrderID();
+          Binding.setCounterImplicitOrderID(OrderID);
+        }
       }
     }
   }
@@ -3808,19 +3855,31 @@ bool SemaHLSL::initGlobalResourceDecl(VarDecl *VD) {
   CXXMethodDecl *CreateMethod = nullptr;
   llvm::SmallVector<Expr *> Args;
 
+  bool HasCounter = hasCounterHandle(ResourceDecl);
+  const char *CreateMethodName;
+  if (Binding.isExplicit())
+    CreateMethodName = HasCounter ? "__createFromBindingWithImplicitCounter"
+                                  : "__createFromBinding";
+  else
+    CreateMethodName = HasCounter
+                           ? "__createFromImplicitBindingWithImplicitCounter"
+                           : "__createFromImplicitBinding";
+
+  CreateMethod =
+      lookupMethod(SemaRef, ResourceDecl, CreateMethodName, VD->getLocation());
+
+  if (!CreateMethod)
+    // This can happen if someone creates a struct that looks like an HLSL
+    // resource record but does not have the required static create method.
+    // No binding will be generated for it.
+    return false;
+
   if (Binding.isExplicit()) {
-    // The resource has explicit binding.
-    CreateMethod = lookupMethod(SemaRef, ResourceDecl, "__createFromBinding",
-                                VD->getLocation());
     IntegerLiteral *RegSlot =
         IntegerLiteral::Create(AST, llvm::APInt(UIntTySize, Binding.getSlot()),
                                AST.UnsignedIntTy, SourceLocation());
     Args.push_back(RegSlot);
   } else {
-    // The resource has implicit binding.
-    CreateMethod =
-        lookupMethod(SemaRef, ResourceDecl, "__createFromImplicitBinding",
-                     VD->getLocation());
     uint32_t OrderID = (Binding.hasImplicitOrderID())
                            ? Binding.getImplicitOrderID()
                            : getNextImplicitBindingOrderID();
@@ -3830,12 +3889,6 @@ bool SemaHLSL::initGlobalResourceDecl(VarDecl *VD) {
     Args.push_back(OrderId);
   }
 
-  if (!CreateMethod)
-    // This can happen if someone creates a struct that looks like an HLSL
-    // resource record but does not have the required static create method.
-    // No binding will be generated for it.
-    return false;
-
   IntegerLiteral *Space =
       IntegerLiteral::Create(AST, llvm::APInt(UIntTySize, Binding.getSpace()),
                              AST.UnsignedIntTy, SourceLocation());
@@ -3859,6 +3912,15 @@ bool SemaHLSL::initGlobalResourceDecl(VarDecl *VD) {
       Name, nullptr, VK_PRValue, FPOptionsOverride());
   Args.push_back(NameCast);
 
+  if (HasCounter) {
+    // Will this be in the correct order?
+    uint32_t CounterOrderID = getNextImplicitBindingOrderID();
+    IntegerLiteral *CounterId =
+        IntegerLiteral::Create(AST, llvm::APInt(UIntTySize, CounterOrderID),
+                               AST.UnsignedIntTy, SourceLocation());
+    Args.push_back(CounterId);
+  }
+
   // Make sure the create method template is instantiated and emitted.
   if (!CreateMethod->isDefined() && CreateMethod->isTemplateInstantiation())
     SemaRef.InstantiateFunctionDefinition(VD->getLocation(), CreateMethod,
@@ -3899,20 +3961,24 @@ bool SemaHLSL::initGlobalResourceArrayDecl(VarDecl *VD) {
   ASTContext &AST = SemaRef.getASTContext();
   QualType ResElementTy = AST.getBaseElementType(VD->getType());
   CXXRecordDecl *ResourceDecl = ResElementTy->getAsCXXRecordDecl();
-
-  HLSLResourceBindingAttr *RBA = VD->getAttr<HLSLResourceBindingAttr>();
-  HLSLVkBindingAttr *VkBinding = VD->getAttr<HLSLVkBindingAttr>();
   CXXMethodDecl *CreateMethod = nullptr;
 
-  if (VkBinding || (RBA && RBA->hasRegisterSlot()))
+  bool HasCounter = hasCounterHandle(ResourceDecl);
+  ResourceBindingAttrs ResourceAttrs(VD);
+  if (ResourceAttrs.isExplicit())
     // Resource has explicit binding.
-    CreateMethod = lookupMethod(SemaRef, ResourceDecl, "__createFromBinding",
-                                VD->getLocation());
-  else
-    // Resource has implicit binding.
     CreateMethod =
-        lookupMethod(SemaRef, ResourceDecl, "__createFromImplicitBinding",
+        lookupMethod(SemaRef, ResourceDecl,
+                     HasCounter ? "__createFromBindingWithImplicitCounter"
+                                : "__createFromBinding",
                      VD->getLocation());
+  else
+    // Resource has implicit binding.
+    CreateMethod = lookupMethod(
+        SemaRef, ResourceDecl,
+        HasCounter ? "__createFromImplicitBindingWithImplicitCounter"
+                   : "__createFromImplicitBinding",
+        VD->getLocation());
 
   if (!CreateMethod)
     return false;