feat: support UseCompressedOops for PEA

llvm/include/llvm/Analysis/Jeandle/PartialEscape.h

@@ -243,6 +243,12 @@ class VirtualObject {
   // callers treat -1 as "keep everything real" (conservative escape).
   int getOrCreateFieldIndex(int64_t Offset, Type *Ty, const DataLayout &DL);
 
+  // Return the tracked field descriptor at an exact byte offset, or nullptr
+  // when the slot has not been observed yet. Materialization uses this to
+  // retain the physical storage type (notably the 32-bit narrow-oop type)
+  // while FieldValue keeps the semantic Java-heap pointer type.
+  const FieldDesc *findField(int64_t Offset) const;
+
   // Result of matching a GEP against the array's element-address pattern.
   // Index is the (possibly symbolic) Value* that names the Java-level
   // element index; ElementType is the per-element LLVM type. Callers must
@@ -543,6 +549,7 @@ class Effect {
     ReplaceCall,
     EliminateStore,
     EliminateAllocation,
+    PlaceInstruction,
     Materialize,
     CreatePHI,
     // Atomically replace one safepoint's complete deoptimization object pool.
@@ -569,8 +576,15 @@ class Effect {
   uint32_t SeqNo = 0;
 
 private:
-  // The one virtual object whose ordinary IR mutation this effect performs.
-  // A deopt-pool rewrite spans zero or more objects and has no mutation owner.
+  // The virtual object whose ordinary IR mutation this effect owns. Ordinary
+  // effects are attached to exactly one VO so commit() can validate their
+  // dependencies and drop/retain them as a unit when that VO becomes
+  // ineligible. InvalidObjectID denotes an effect that is deliberately
+  // independent of any one VO: RewriteDeoptPool rewrites one complete
+  // safepoint pool and may cover several (or no) VOs, while
+  // PlaceInstruction only parents an analysis-created instruction at its
+  // program point. Ownerless effects must therefore not participate in the
+  // per-VO ineligibility cascade.
   ObjectID MutationOwner = InvalidObjectID;
 
 public:
@@ -592,8 +606,10 @@ class Effect {
   virtual Kind getKind() const = 0;
 
   bool hasValidMutationOwner() const {
-    return getKind() == Kind::RewriteDeoptPool ? !hasMutationOwner()
-                                               : hasMutationOwner();
+    return (getKind() == Kind::RewriteDeoptPool ||
+            getKind() == Kind::PlaceInstruction)
+               ? !hasMutationOwner()
+               : hasMutationOwner();
   }
 
   // The IR instruction this effect rewrites/erases, or null for effects that
@@ -683,6 +699,27 @@ class EliminateStoreEffect : public Effect {
   }
 };
 
+// Parent an analysis-created instruction at its modeled program point.
+class PlaceInstructionEffect : public Effect {
+public:
+  // Both handles are weak: an earlier ordinary effect may erase the modeled
+  // store before this placement runs.
+  WeakTrackingVH Target;
+  WeakTrackingVH InstructionToPlace;
+
+  Kind getKind() const override { return Kind::PlaceInstruction; }
+  static bool classof(const Effect *E) {
+    return E->getKind() == Kind::PlaceInstruction;
+  }
+  Instruction *getTarget() const override {
+    return dyn_cast_or_null<Instruction>((Value *)Target);
+  }
+  void apply(TransformContext &Ctx) override;
+  std::unique_ptr<Effect> clone() const override {
+    return std::make_unique<PlaceInstructionEffect>(*this);
+  }
+};
+
 // Rewrite the original allocation invoke into an unconditional branch (dropping
 // the unwind edge) or erase a call alloc. Applied in the cfg-kill phase.
 class EliminateAllocationEffect : public Effect {
@@ -719,8 +756,15 @@ class MaterializeEffect : public Effect {
   // Per-offset snapshot of a virtual object's field values at a
   // materialization point.
   struct FieldEntry {
-    int64_t Offset;
+    // Physical slot description. Value's declared type is the semantic Java
+    // value (AS1 for references), whereas compressed-oop fields are physically
+    // stored as AS3 pointers.
+    VirtualObject::FieldDesc Storage;
     FieldValue Value;
+
+    FieldEntry() : Storage{0, nullptr, 0, false} {}
+    FieldEntry(const VirtualObject::FieldDesc &FD, const FieldValue &FV)
+        : Storage(FD), Value(FV) {}
   };
 
   // WeakTrackingVH so erasing the insertion-point instruction auto-nulls the

llvm/lib/Analysis/Jeandle/PartialEscape.cpp

@@ -81,17 +81,14 @@ int VirtualObject::getOrCreateFieldIndex(int64_t Offset, Type *Ty,
     // current 64-bit target, but derived from the DataLayout so a 32-bit or
     // compressed-oop heap model stays correct rather than hardcoding 8.
     //
-    // TODO(compressed-oop): narrow-oop (addrspace 3) reference fields are NOT
-    // supported — bail conservatively (-1) instead of asserting (debug) or
-    // modelling the slot at the wrong width (release: getPointerSize(1)=8
-    // where the real slot is 4 bytes -> corrupt field model). Callers treat
-    // -1 as keep-everything-real. PEA as a whole is also gated against
-    // narrow-oop modules in PartialEscapeAnalysis::run; this is the
-    // per-access defense for hand-written / mixed IR.
-    if (Ty->getPointerAddressSpace() != jeandle::AddrSpace::JavaHeapAddrSpace)
+    // Reference fields may use the semantic Java heap pointer (AS1) or the
+    // physical narrow-oop pointer (AS3). Keep the descriptor's type so
+    // materialization can replay the value at its actual in-memory width.
+    unsigned PointerAS = Ty->getPointerAddressSpace();
+    if (PointerAS != jeandle::AddrSpace::JavaHeapAddrSpace &&
+        PointerAS != jeandle::AddrSpace::NarrowOopAddrSpace)
       return -1;
-    uint64_t PointerByteSize =
-        DL.getPointerSize(jeandle::AddrSpace::JavaHeapAddrSpace);
+    uint64_t PointerByteSize = DL.getPointerSize(PointerAS);
     if (PointerByteSize == 0 ||
         PointerByteSize > std::numeric_limits<uint8_t>::max())
       return -1;
@@ -150,6 +147,14 @@ int VirtualObject::getOrCreateFieldIndex(int64_t Offset, Type *Ty,
   auto NewIt = Fields.insert(It, New);
   return static_cast<int>(NewIt - Fields.begin());
 }
+const VirtualObject::FieldDesc *VirtualObject::findField(int64_t Offset) const {
+  auto It = std::lower_bound(
+      Fields.begin(), Fields.end(), Offset,
+      [](const FieldDesc &F, int64_t Off) { return F.Offset < Off; });
+  if (It == Fields.end() || It->Offset != Offset)
+    return nullptr;
+  return &*It;
+}
 
 // Strip identity-preserving wrappers (freeze, bitcast, zext, sext) from an
 // index Value. Used by matchArrayElementGEP to canonicalize the index so
@@ -386,9 +391,11 @@ FieldValue FieldValue::materializedRef(Value *Ptr) {
 Constant *FieldValue::defaultFor(Type *FieldType) {
   assert(FieldType);
   if (FieldType->isPointerTy()) {
-    assert(FieldType->getPointerAddressSpace() ==
-               jeandle::AddrSpace::JavaHeapAddrSpace &&
-           "reference default must be in JavaHeapAddrSpace");
+    assert((FieldType->getPointerAddressSpace() ==
+                jeandle::AddrSpace::JavaHeapAddrSpace ||
+            FieldType->getPointerAddressSpace() ==
+                jeandle::AddrSpace::NarrowOopAddrSpace) &&
+           "reference default must be a Java oop pointer");
     return ConstantPointerNull::get(cast<PointerType>(FieldType));
   }
   return Constant::getNullValue(FieldType);
@@ -409,10 +416,6 @@ bool FieldValue::shallowEquals(const FieldValue &O) const {
   return false;
 }
 
-// ===========================================================================
-// ObjectState
-// ===========================================================================
-
 // ===========================================================================
 // PEABlockState
 // ===========================================================================
@@ -494,8 +497,12 @@ ObjectState &PEABlockState::getObjectStateForModification(ObjectID ID) {
 
 void AliasMap::addVirtualAlias(Value *V, ObjectID ID, bool IsWholeObject) {
   assert(V && ID != InvalidObjectID);
-  assert(!VirtualAliases.count(V) && "value already aliased");
-  VirtualAliases[V] = ID;
+  auto It = VirtualAliases.find(V);
+  if (It != VirtualAliases.end()) {
+    assert(It->second == ID && "value already aliased to a different object");
+  } else {
+    VirtualAliases[V] = ID;
+  }
   if (IsWholeObject)
     WholeObjectVirtualAliases.insert(V);
   for (User *U : V->users()) {
@@ -770,7 +777,7 @@ ObjectID PEAResult::createVirtualObject(std::unique_ptr<VirtualObject> VO) {
 
 void PEAResult::addBlockEffect(std::unique_ptr<Effect> E) {
   assert(E && E->hasValidMutationOwner() &&
-         "only an atomic deopt-pool effect may be ownerless");
+         "only deopt-pool or placement effects may be ownerless");
   assert(E->Block);
   BasicBlock *BB = E->Block;
   BlockEffects[BB].add(std::move(E));
@@ -825,6 +832,9 @@ void Effect::dump(raw_ostream &OS) const {
   case Kind::EliminateStore:
     OS << "EliminateStore";
     break;
+  case Kind::PlaceInstruction:
+    OS << "PlaceInstruction";
+    break;
   case Kind::EliminateAllocation:
     OS << "EliminateAllocation";
     break;