[llvm] [llubi] Add support for exposed provenance (PR #200596)

[Nikita Popov via llvm-commits]

================
@@ -612,6 +615,103 @@ Pointer Context::deriveFromMemoryObject(IntrusiveRefCntPtr<MemoryObject> Obj) {
                        Obj->getAddress()));
 }
 
+void Context::exposeProvenance(Provenance &Prov) {
+  if (Prov.Wildcard)
+    return;
+  MemoryObject *Obj = Prov.getMemoryObject();
+  if (!Obj)
+    return;
+  uint64_t Address = Obj->getAddress();
+  ExposedProvenanceSet &Set = ExposedProvenances[Address];
+  if (Set.Set.insert(&Prov).second)
+    Set.List.push_back({&Prov, ++ExposedProvenanceSetGeneration});
+}
+
+MemoryObject *
+Context::checkProvenance(const Pointer &Ptr,
+                         function_ref<bool(const Provenance &)> Check,
+                         bool HasSideEffect) {
+  auto &Prov = Ptr.provenance();
+  if (!Check(Prov))
+    return nullptr;
+  // Early return for concrete provenances.
+  if (!Prov.Wildcard)
+    return Prov.Obj.get();
+
+  MemoryObject *MO = nullptr;
+  APInt &Mask = Prov.Wildcard->ActiveMask;
+  SmallVector<ExposedProvenance> *List = nullptr;
+  uint32_t ProvenanceCount = 0;
+  if (Mask.isZero()) {
+    // The memory object hasn't been determined.
+    uint64_t Addr = Ptr.address().getLimitedValue();
+    auto Iter = ExposedProvenances.upper_bound(Addr);
+    if (Iter == ExposedProvenances.begin())
+      return nullptr;
+    auto &[BaseAddress, Set] = *std::prev(Iter);
+    auto &Obj = MemoryObjects.at(BaseAddress);
+    if (!Obj->inBounds(Ptr.address()))
+      return nullptr;
+    MO = Obj.get();
+    // We only inspect the first N exposed provenances according to the global
+    // generation number of the wildcard pointer.
+    ProvenanceCount = std::distance(
+        Set.List.begin(),
+        upper_bound(Set.List,
+                    ExposedProvenance{nullptr, Prov.Wildcard->Generation}));
+    if (HasSideEffect) {
+      Mask = APInt::getAllOnes(ProvenanceCount);
+      Prov.Wildcard->BaseAddress = BaseAddress;
+    }
+    List = &Set.List;
+  } else {
+    // We already determined the memory object in a previous memory access.
+    uint64_t BaseAddress = Prov.Wildcard->BaseAddress;
+    auto Iter = ExposedProvenances.find(BaseAddress);
+    // The memory object has been freed.
+    if (Iter == ExposedProvenances.end())
+      return nullptr;
----------------
nikic wrote:

Missing test coverage for this (wildcard provenance of freed object).

https://github.com/llvm/llvm-project/pull/200596