New uniform mapping handling

- add optional callback to handle mapping of uniform variables in linking phase
- if no resolver is provided, it uses the internal default resolver with all shifts and auto bind settings

Change-Id: Icfe38a9eabe8bfc8f8bb6d8150c06f7ed38bb762

5 files changed, 297 insertions, 156 deletions

diff --git a/Test/baseResults/spv.register.autoassign.rangetest.frag.out b/Test/baseResults/spv.register.autoassign.rangetest.frag.out
index a521a13b..94933fba 100644
--- a/Test/baseResults/spv.register.autoassign.rangetest.frag.out
+++ b/Test/baseResults/spv.register.autoassign.rangetest.frag.out
@@ -2,10 +2,6 @@ spv.register.autoassign.rangetest.frag
 
 Linked fragment stage:
 
-INTERNAL ERROR: mapped binding out of range: g_tScene
-INTERNAL ERROR: mapped binding out of range: g_tSamp
-INTERNAL ERROR: mapped binding out of range: g_tScene
-INTERNAL ERROR: mapped binding out of range: g_tSamp
 INTERNAL ERROR: mapped binding out of range: g_tSamp
 INTERNAL ERROR: mapped binding out of range: g_tScene
 
diff --git a/glslang/MachineIndependent/ShaderLang.cpp b/glslang/MachineIndependent/ShaderLang.cpp
index 6d044452..eb4a17d4 100644
--- a/glslang/MachineIndependent/ShaderLang.cpp
+++ b/glslang/MachineIndependent/ShaderLang.cpp
@@ -1716,7 +1716,7 @@ void TProgram::dumpReflection()                      { reflection->dump(); }
 //
 // I/O mapping implementation.
 //
-bool TProgram::mapIO()
+bool TProgram::mapIO(TIoMapResolver* resolver)
 {
     if (! linked || ioMapper)
         return false;
@@ -1725,7 +1725,7 @@ bool TProgram::mapIO()
 
     for (int s = 0; s < EShLangCount; ++s) {
         if (intermediate[s]) {
-            if (! ioMapper->addStage((EShLanguage)s, *intermediate[s], *infoSink))
+            if (! ioMapper->addStage((EShLanguage)s, *intermediate[s], *infoSink, resolver))
                 return false;
         }
     }
diff --git a/glslang/MachineIndependent/iomapper.cpp b/glslang/MachineIndependent/iomapper.cpp
index 15847ddb..8ea4506f 100644
--- a/glslang/MachineIndependent/iomapper.cpp
+++ b/glslang/MachineIndependent/iomapper.cpp
@@ -61,177 +61,278 @@
 //    c. implicit dead bindings are left un-bound.
 //
 
-   
-namespace glslang {
 
-// Map of IDs to bindings
-typedef std::unordered_map<unsigned int, int> TBindingMap;
-typedef std::unordered_set<int> TUsedBindings;
+namespace glslang {
 
+struct TVarEntryInfo
+{
+    int             id;
+    TIntermSymbol*  symbol;
+    bool            live;
+    int             newBinding;
+    int             newSet;
 
-// This traverses the AST to determine which bindings are used, and which are implicit
-// (for subsequent auto-numbering)
-class TBindingTraverser : public TLiveTraverser {
-public:
-    TBindingTraverser(const TIntermediate& i, TBindingMap& bindingMap, TUsedBindings& usedBindings,
-                      bool traverseDeadCode = false) :
-        TLiveTraverser(i, traverseDeadCode, true, true, false),
-        bindingMap(bindingMap),
-        usedBindings(usedBindings)
-    { }
-
-protected:
-    virtual void visitSymbol(TIntermSymbol* base) {
-        if (base->getQualifier().storage == EvqUniform)
-            addUniform(*base);
-    }
+    struct TOrderById
+    {
+      inline bool operator()(const TVarEntryInfo& l, const TVarEntryInfo& r)
+      {
+        return l.id < r.id;
+      }
+    };
 
-    // Return the right binding base given the variable type.
-    int getBindingBase(const TType& type) {
-        if (type.getBasicType() == EbtSampler) {
-            const TSampler& sampler = type.getSampler();
-            if (sampler.isPureSampler())
-                return intermediate.getShiftSamplerBinding();
-            if (sampler.isTexture())
-                return intermediate.getShiftTextureBinding();
+    struct TOrderByPriority
+    {
+        // ordering:
+        // 1) has both binding and set
+        // 2) has binding but no set
+        // 3) has no binding but set
+        // 4) has no binding and no set
+        inline bool operator()(const TVarEntryInfo& l, const TVarEntryInfo& r)
+        {
+            const TQualifier& lq = l.symbol->getQualifier();
+            const TQualifier& rq = r.symbol->getQualifier();
+
+            // simple rules:
+            // has binding gives 2 points
+            // has set gives 1 point
+            // who has the most points is more important.
+            int lPoints = (lq.hasBinding() ? 2 : 0) + (lq.hasSet() ? 1 : 0);
+            int rPoints = (rq.hasBinding() ? 2 : 0) + (rq.hasSet() ? 1 : 0);
+
+            if (lPoints == rPoints)
+              return l.id < r.id;
+            return lPoints > rPoints;
         }
+    };
+};
 
-        if (type.getQualifier().isUniformOrBuffer())
-            return intermediate.getShiftUboBinding();
-
-        return -1;  // not a type with a binding
-    }
 
-    // Mark a given base symbol ID as being bound to 'binding'
-    void markBinding(const TIntermSymbol& base, int binding) {
-        bindingMap[base.getId()] = binding;
 
-        if (binding >= 0) {
-            // const TType& type = base.getType();
-            const unsigned int size = 1; // type.isArray() ? type.getCumulativeArraySize() : 1;
+typedef std::vector<TVarEntryInfo> TVarLiveMap;
 
-            for (unsigned int offset=0; offset<size; ++offset)
-                usedBindings.insert(binding + offset);
-        }
-    }
-
-    // Mark the bindings that are given explicitly, and set ones that need
-    // implicit bindings to -1 for a subsequent pass.  (Can't happen in this
-    // pass because explicit bindings in dead code reserve the location).
-    virtual void addUniform(TIntermSymbol& base)
+class TVarGatherTraverser : public TLiveTraverser
+{
+public:
+    TVarGatherTraverser(const TIntermediate& i, TVarLiveMap& vars, bool traverseDeadCode)
+      : TLiveTraverser(i, traverseDeadCode, true, true, false)
+      , varLiveList(vars)
     {
-        // Skip ones we've already seen.
-        if (bindingMap.find(base.getId()) != bindingMap.end())
-            return;
+    }
 
-        const TType& type = base.getType();
-        const int bindingBase = getBindingBase(type);
 
-        // Return if it's not a type we bind
-        if (bindingBase == -1)
-            return;
-
-        if (type.getQualifier().hasBinding()) {
-            // It has a binding: keep that one.
-            markBinding(base, type.getQualifier().layoutBinding + bindingBase);
-        } else if (!traverseAll) {
-            // Mark it as something we need to dynamically create a binding for,
-            // only if we're walking just the live code.  We don't auto-number
-            // in dead code.
-            markBinding(base, -1);
+    virtual void visitSymbol(TIntermSymbol* base)
+    {
+        if (base->getQualifier().storage == EvqUniform) {
+            TVarEntryInfo ent = { base->getId(), base, !traverseAll };
+            TVarLiveMap::iterator at = std::lower_bound(varLiveList.begin(), varLiveList.end(), ent, TVarEntryInfo::TOrderById());
+            if (at != varLiveList.end() && at->id == ent.id)
+              at->live = at->live || !traverseAll; // update live state
+            else
+              varLiveList.insert(at, ent);
         }
     }
 
-    TBindingMap&         bindingMap;
-    TUsedBindings&       usedBindings;
+  private:
+    TVarLiveMap&    varLiveList;
 };
 
-
-// This traverses the AST and applies binding maps it's given.
-class TIoMappingTraverser : public TBindingTraverser {
+class TVarSetTraverser : public TLiveTraverser
+{
 public:
-    TIoMappingTraverser(TIntermediate& i, TBindingMap& bindingMap, TUsedBindings& usedBindings,
-                        TInfoSink& infoSink, bool traverseDeadCode) :
-        TBindingTraverser(i, bindingMap, usedBindings, traverseDeadCode),
-        infoSink(infoSink),
-        assignError(false)
-    { }
-
-    bool success() const { return !assignError; }
-
-protected:
-    unsigned checkBindingRange(const TIntermSymbol& base, unsigned binding)
+    TVarSetTraverser(const TIntermediate& i, const TVarLiveMap& vars)
+      : TLiveTraverser(i, true, true, true, false)
+      , varLiveList(vars)
     {
-        if (binding >= TQualifier::layoutBindingEnd) {
-            TString err = "mapped binding out of range: ";
-            err += base.getName();
-
-            infoSink.info.message(EPrefixInternalError, err.c_str());
-            assignError = true;
-            
-            return 0;
-        }
-
-        return binding;
     }
 
-    void addUniform(TIntermSymbol& base) override
+
+    virtual void visitSymbol(TIntermSymbol* base)
     {
-        // Skip things we don't intend to bind.
-        if (bindingMap.find(base.getId()) == bindingMap.end())
+        TVarEntryInfo ent = { base->getId() };
+        TVarLiveMap::const_iterator at = std::lower_bound(varLiveList.begin(), varLiveList.end(), ent, TVarEntryInfo::TOrderById());
+        if (at == varLiveList.end())
             return;
-
-        const int existingBinding = bindingMap[base.getId()];
-
-        // Apply existing binding, if we were given one or already made one up.
-        if (existingBinding != -1) {
-            base.getWritableType().getQualifier().layoutBinding = checkBindingRange(base, existingBinding);
+        if (!(at->id == ent.id))
             return;
-        }
 
-        if (intermediate.getAutoMapBindings()) {
-            // Otherwise, find a free spot for it.
-            const int freeBinding = getFreeBinding(base.getType(), getBindingBase(base.getType()));
+        if (at->newBinding != -1)
+            base->getWritableType().getQualifier().layoutBinding = at->newBinding;
+        if (at->newSet != -1)
+            base->getWritableType().getQualifier().layoutSet = at->newSet;
+    }
+
+  private:
+    const TVarLiveMap&    varLiveList;
+};
 
-            markBinding(base, freeBinding);
-            base.getWritableType().getQualifier().layoutBinding = checkBindingRange(base, freeBinding);
+struct TResolverAdaptor
+{
+    TResolverAdaptor(EShLanguage s, TIoMapResolver& r, TInfoSink& i, bool& e)
+      : resolver(r)
+      , stage(s)
+      , infoSink(i)
+      , error(e)
+    {
+    }
+    inline void operator()(TVarEntryInfo& ent)
+    {
+        bool isValid = resolver.validateBinding(stage, ent.symbol->getName().c_str(), ent.symbol->getType(), ent.live);
+        if (isValid) {
+            ent.newBinding = resolver.resolveBinding(stage, ent.symbol->getName().c_str(), ent.symbol->getType(), ent.live);
+            ent.newSet = resolver.resolveSet(stage, ent.symbol->getName().c_str(), ent.symbol->getType(), ent.live);
+
+            if (ent.newBinding != -1) {
+                if (ent.newBinding >= TQualifier::layoutBindingEnd) {
+                    TString err = "mapped binding out of range: " + ent.symbol->getName();
+
+                    infoSink.info.message(EPrefixInternalError, err.c_str());
+                    error = true;
+                }
+            }
+            if (ent.newSet != -1) {
+                if (ent.newSet >= TQualifier::layoutSetEnd) {
+                    TString err = "mapped set out of range: " + ent.symbol->getName();
+
+                    infoSink.info.message(EPrefixInternalError, err.c_str());
+                    error = true;
+                }
+            }
+        } else {
+            TString errorMsg = "Invalid binding: " + ent.symbol->getName();
+            infoSink.info.message(EPrefixInternalError, errorMsg.c_str());
+            error = true;
         }
     }
+    EShLanguage     stage;
+    TIoMapResolver& resolver;
+    TInfoSink&      infoSink;
+    bool&           error;
+};
 
-    // Search for N free consecutive binding slots in [base, base+required).
-    // E.g, if we want to reserve consecutive bindings for flattened arrays.
-    bool hasNFreeSlots(int base, int required) {
-        for (int binding = base; binding < (base + required); ++binding)
-            if (usedBindings.find(binding) != usedBindings.end())
-                return false;
-        
-        return true;
+/*
+ * Basic implementation of glslang::TIoMapResolver that replaces the
+ * previous offset behaviour.
+ * It does the same, uses the offsets for th corresponding uniform
+ * types. Also respects the EOptionAutoMapBindings flag and binds
+ * them if needed.
+ */
+struct TDefaultIoResolver : public glslang::TIoMapResolver
+{
+  int baseSamplerBinding;
+  int baseTextureBinding;
+  int baseUboBinding;
+  bool doAutoMapping;
+  typedef std::vector<int> TSlotSet;
+  typedef std::unordered_map<int, TSlotSet> TSlotSetMap;
+  TSlotSetMap slots;
+  TSlotSet::iterator findSlot(int set, int slot)
+  {
+    return std::lower_bound(slots[set].begin(), slots[set].end(), slot);
+  }
+  bool checkEmpty(int set, int slot)
+  {
+    TSlotSet::iterator at = findSlot(set, slot);
+    return !(at != slots[set].end() && *at == slot);
+  }
+  int reserveSlot(int set, int slot)
+  {
+    TSlotSet::iterator at = findSlot(set, slot);
+    slots[set].insert(at, slot);
+    return slot;
+  }
+  int getFreeSlot(int set, int base)
+  {
+    TSlotSet::iterator at = findSlot(set, base);
+    if (at == slots[set].end())
+      return reserveSlot(set, base);
+
+    // look in locksteps, if they not match, then there is a free slot
+    for (; at != slots[set].end(); ++at, ++base)
+      if (*at != base)
+        break;
+    return reserveSlot(set, base);
+  }
+  bool validateBinding(EShLanguage stage, const char* /*name*/, const glslang::TType& type, bool /*is_live*/) override
+  {
+    if (type.getQualifier().hasBinding()) {
+      int set;
+      if (type.getQualifier().hasSet())
+        set = type.getQualifier().layoutSet;
+      else
+        set = 0;
+
+      if (type.getBasicType() == glslang::EbtSampler) {
+        const glslang::TSampler& sampler = type.getSampler();
+        if (sampler.isPureSampler())
+          return checkEmpty(set, baseSamplerBinding + type.getQualifier().layoutBinding);
+
+        if (sampler.isTexture())
+          return checkEmpty(set, baseTextureBinding + type.getQualifier().layoutBinding);
+      }
+
+      if (type.getQualifier().isUniformOrBuffer())
+        return checkEmpty(set, baseUboBinding + type.getQualifier().layoutBinding);
     }
-
-    // Find a free binding spot 
-    int getFreeBinding(const TType&, int nextBinding) {
-        while (!hasNFreeSlots(nextBinding, 1))
-            ++nextBinding;
-
-        return nextBinding;
+    return true;
+  }
+  int resolveBinding(EShLanguage stage, const char* /*name*/, const glslang::TType& type, bool is_live) override
+  {
+    int set;
+    if (type.getQualifier().hasSet())
+      set = type.getQualifier().layoutSet;
+    else
+      set = 0;
+
+    if (type.getQualifier().hasBinding()) {
+      if (type.getBasicType() == glslang::EbtSampler) {
+        const glslang::TSampler& sampler = type.getSampler();
+        if (sampler.isPureSampler())
+          return reserveSlot(set, baseSamplerBinding + type.getQualifier().layoutBinding);
+
+        if (sampler.isTexture())
+          return reserveSlot(set, baseTextureBinding + type.getQualifier().layoutBinding);
+      }
+
+      if (type.getQualifier().isUniformOrBuffer())
+        return reserveSlot(set, baseUboBinding + type.getQualifier().layoutBinding);
+    } else if (is_live && doAutoMapping) {
+      // find free slot, the caller did make sure it passes all vars with binding
+        // first and now all are passed that do not have a binding and needs one
+      if (type.getBasicType() == glslang::EbtSampler) {
+        const glslang::TSampler& sampler = type.getSampler();
+        if (sampler.isPureSampler())
+          return getFreeSlot(set, baseSamplerBinding);
+
+        if (sampler.isTexture())
+          return getFreeSlot(set, baseTextureBinding);
+      }
+
+      if (type.getQualifier().isUniformOrBuffer())
+        return getFreeSlot(set, baseUboBinding);
     }
 
-private:
-    bool assignError;  // true if there was an error assigning the bindings
-    TInfoSink& infoSink;
+    return -1;
+  }
+  int resolveSet(EShLanguage /*stage*/, const char* /*name*/, const glslang::TType& type, bool /*is_live*/) override
+  {
+    if (type.getQualifier().hasSet())
+      return type.getQualifier().layoutSet;
+    return 0;
+  }
 };
 
 // Map I/O variables to provided offsets, and make bindings for
 // unbound but live variables.
 //
 // Returns false if the input is too malformed to do this.
-bool TIoMapper::addStage(EShLanguage, TIntermediate& intermediate, TInfoSink& infoSink)
+bool TIoMapper::addStage(EShLanguage stage, TIntermediate &intermediate, TInfoSink &infoSink, TIoMapResolver *resolver)
 {
     // Trivial return if there is nothing to do.
     if (intermediate.getShiftSamplerBinding() == 0 &&
         intermediate.getShiftTextureBinding() == 0 &&
         intermediate.getShiftUboBinding() == 0 &&
-        intermediate.getAutoMapBindings() == false)
+        intermediate.getAutoMapBindings() == false &&
+        resolver == NULL)
         return true;
 
     if (intermediate.getNumEntryPoints() != 1 || intermediate.isRecursive())
@@ -241,30 +342,43 @@ bool TIoMapper::addStage(EShLanguage, TIntermediate& intermediate, TInfoSink& in
     if (root == nullptr)
         return false;
 
-    // The lifetime of this data spans several passes.
-    TBindingMap   bindingMap;
-    TUsedBindings usedBindings;
+    // if no resolver is provided, use the default resolver with the given shifts and auto map settings
+    TDefaultIoResolver defaultResolver;
+    if (resolver == NULL) {
+      defaultResolver.baseSamplerBinding = intermediate.getShiftSamplerBinding();
+      defaultResolver.baseTextureBinding = intermediate.getShiftTextureBinding();
+      defaultResolver.baseUboBinding = intermediate.getShiftUboBinding();
+      defaultResolver.doAutoMapping = intermediate.getAutoMapBindings();
 
-    TBindingTraverser it_binding_all(intermediate, bindingMap, usedBindings, true);
-    TBindingTraverser it_binding_live(intermediate, bindingMap, usedBindings, false);
-    TIoMappingTraverser it_iomap(intermediate, bindingMap, usedBindings, infoSink, true);
+      resolver = &defaultResolver;
+    }
 
-    // Traverse all (live+dead) code to find explicit bindings, so we can avoid those.
-    root->traverse(&it_binding_all);
+    TVarLiveMap varMap;
+    TVarGatherTraverser iter_binding_all(intermediate, varMap, true);
+    TVarGatherTraverser iter_binding_live(intermediate, varMap, false);
 
-    // Traverse just live code to find things that need implicit bindings.
-    it_binding_live.pushFunction(intermediate.getEntryPointMangledName().c_str());
+    root->traverse(&iter_binding_all);
+    iter_binding_live.pushFunction(intermediate.getEntryPointMangledName().c_str());
 
-    while (! it_binding_live.functions.empty()) {
-        TIntermNode* function = it_binding_live.functions.back();
-        it_binding_live.functions.pop_back();
-        function->traverse(&it_binding_live);
+    while (!iter_binding_live.functions.empty()) {
+        TIntermNode* function = iter_binding_live.functions.back();
+        iter_binding_live.functions.pop_back();
+        function->traverse(&iter_binding_live);
     }
-
-    // Bind everything that needs a binding and doesn't have one.
-    root->traverse(&it_iomap);
-
-    return it_iomap.success();
+    // sort entries by priority. see TVarEntryInfo::TOrderByPriority for info.
+    std::sort(varMap.begin(), varMap.end(), TVarEntryInfo::TOrderByPriority());
+
+    bool hadError = false;
+    TResolverAdaptor doResolve(stage, *resolver, infoSink, hadError);
+    std::for_each(varMap.begin(), varMap.end(), doResolve);
+
+    if (!hadError) {
+      // sort by id again, so we can use lower bound to find entries
+      std::sort(varMap.begin(), varMap.end(), TVarEntryInfo::TOrderById());
+      TVarSetTraverser iter_iomap(intermediate, varMap);
+      root->traverse(&iter_iomap);
+    }
+    return !hadError;
 }
 
 } // end namespace glslang
diff --git a/glslang/MachineIndependent/iomapper.h b/glslang/MachineIndependent/iomapper.h
index 68dec776..ff47543d 100644
--- a/glslang/MachineIndependent/iomapper.h
+++ b/glslang/MachineIndependent/iomapper.h
@@ -55,7 +55,7 @@ public:
     virtual ~TIoMapper() {}
 
     // grow the reflection stage by stage
-    bool addStage(EShLanguage, TIntermediate&, TInfoSink&);
+    bool addStage(EShLanguage, TIntermediate&, TInfoSink&, TIoMapResolver*);
 };
 
 } // end namespace glslang
diff --git a/glslang/Public/ShaderLang.h b/glslang/Public/ShaderLang.h
index e5e8b4d7..acde8b01 100644
--- a/glslang/Public/ShaderLang.h
+++ b/glslang/Public/ShaderLang.h
@@ -445,7 +445,36 @@ private:
 class TReflection;
 class TIoMapper;
 
-// Make one TProgram per set of shaders that will get linked together.  Add all 
+// Allows to customize the binding layout after linking.
+// All used uniform variables will invoke at least validateBinding.
+// If validateBinding returned true then the other resolveBinding
+// and resolveSet are invoked to resolve the binding and descriptor
+// set index respectively.
+// Invocations happen in a particular order:
+// 1) var with binding and set already defined
+// 2) var with binding but no set defined
+// 3) var with set but no binding defined
+// 4) var with no binding and no set defined
+//
+// NOTE: that still limit checks are applied to bindings and sets
+// and may result in an error.
+class TIoMapResolver
+{
+public:
+  virtual ~TIoMapResolver() {}
+
+  // Should return true if the resulting/current binding would be ok.
+  // Basic idea is to do aliasing binding checks with this.
+  virtual bool validateBinding(EShLanguage stage, const char* name, const TType& type, bool is_live) = 0;
+  // Should return a value >= 0 if the current binding should be overridden.
+  // Return -1 if the current binding (including no binding) should be kept.
+  virtual int resolveBinding(EShLanguage stage, const char* name, const TType& type, bool is_live) = 0;
+  // Should return a value >= 0 if the current set should be overriden.
+  // Return -1 if the current set (including no set) should be kept.
+  virtual int resolveSet(EShLanguage stage, const char* name, const TType& type, bool is_live) = 0;
+};
+
+// Make one TProgram per set of shaders that will get linked together.  Add all
 // the shaders that are to be linked together.  After calling shader.parse()
 // for all shaders, call link().
 //
@@ -485,7 +514,9 @@ public:
     void dumpReflection();
 
     // I/O mapping: apply base offsets and map live unbound variables
-    bool mapIO();
+    // If resolver is not provided it uses the previous approach
+    // and respects auto assignment and offsets.
+    bool mapIO(TIoMapResolver* resolver = NULL);
 
 protected:
     bool linkStage(EShLanguage, EShMessages);