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

[Yingwei Zheng via llvm-commits]

================
@@ -612,6 +615,96 @@ 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.emplace_back(&Prov);
+    Set.GenerationList.push_back(++ExposedProvenanceSetGeneration);
+  }
+}
+
+MemoryObject *
+Context::checkProvenance(const Pointer &Ptr,
+                         function_ref<bool(const Provenance &)> Check,
+                         unsigned AS, 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 TempMask;
+  APInt *Mask = HasSideEffect ? &Prov.Wildcard->ActiveMask : &TempMask;
+  SmallVector<IntrusiveRefCntPtr<Provenance>> *List = nullptr;
+  if (Prov.Wildcard->ActiveMask.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->getAddressSpace() != AS || !Obj->inBounds(Ptr.address()))
+      return nullptr;
+    MO = Obj.get();
+    uint32_t Generation = std::distance(
+        Set.GenerationList.begin(),
+        upper_bound(Set.GenerationList, Prov.Wildcard->Generation));
+    *Mask = APInt::getAllOnes(Generation);
+    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;
+    MO = MemoryObjects.at(BaseAddress).get();
+    if (MO->getAddressSpace() != AS || !MO->inBounds(Ptr.address()))
+      return nullptr;
+    if (HasSideEffect)
+      TempMask = Prov.Wildcard->ActiveMask;
+    List = &Iter->second.List;
+  }
+  if (Prov.Obj) {
+    // We already determined the memory object via speculatable operations like
+    // gep inbounds.
+    if (Prov.Obj.get() != MO)
+      return nullptr;
+  }
+
+  uint32_t Generation = Mask->getBitWidth();
+  bool Valid = false;
+  for (uint32_t I = 0; I != Generation; ++I) {
+    if (!(*Mask)[I])
+      continue;
+    if (Check(*(*List)[I]))
+      Valid = true;
+    else
+      Mask->clearBit(I);
+  }
----------------
dtcxzyw wrote:

See my previous comment https://github.com/llvm/llvm-project/pull/200596#discussion_r3408044267. It is associated with a specific memory object. The actual number of exposed provenance is small in practice if we implemented the FIXME in `ExposedProvenanceSet`. The provenance with `read_provenance` can be ignored if the one with `provenance` has been exposed.


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