[llvm-branch-commits] [clang] [LifetimeSafety][NFC] Add field-labeled child edges to OriginNode and generalize subtree walks (PR #201510)

[Utkarsh Saxena via llvm-branch-commits]

================
@@ -66,47 +66,88 @@ struct Origin {
   }
 };
 
-/// A list of origins representing levels of indirection for pointer-like types.
+/// A tree of origins representing the structure of a pointer-like or
+/// record type.
 ///
-/// Each node in the list contains an OriginID representing a level of
-/// indirection. The list structure captures the multi-level nature of
-/// pointer and reference types in the lifetime analysis.
+/// Each node carries an OriginID and is connected to children via labeled
+/// edges: either a pointee edge (one level of pointer/reference indirection)
+/// or a field edge (a named field of a record). Pointer-like types form a
+/// pointee chain; record types fan out via field edges.
 ///
 /// Examples:
-///   - For `int& x`, the list has size 2:
-///     * Outer: origin for the reference storage itself (the lvalue `x`)
-///     * Inner: origin for what `x` refers to
+///   - For `int& x`, the chain has length 2:
+///       Outer: the reference storage itself (the lvalue `x`)
+///        +-pointee-> Inner: what `x` refers to
 ///
-///   - For `int* p`, the list has size 2:
-///     * Outer: origin for the pointer variable `p`
-///     * Inner: origin for what `p` points to
+///   - For `int* p`, the chain has length 2:
+///       Outer: the pointer variable `p`
+///        +-pointee-> Inner: what `p` points to
 ///
-///   - For `View v` (where View is gsl::Pointer), the list has size 2:
-///     * Outer: origin for the view object itself
-///     * Inner: origin for what the view refers to
+///   - For `View v` (View is gsl::Pointer), the chain has length 2:
+///       Outer: the view object itself
+///        +-pointee-> Inner: what the view refers to
 ///
-///   - For `int** pp`, the list has size 3:
-///     * Outer: origin for `pp` itself
-///     * Inner: origin for `*pp` (what `pp` points to)
-///     * Inner->Inner: origin for `**pp` (what `*pp` points to)
+///   - For `int** pp`, the chain has length 3:
+///       Outer: `pp` itself
+///        +-pointee-> Inner: `*pp` (what `pp` points to)
+///           +-pointee-> Inner->Inner: `**pp` (what `*pp` points to)
 ///
-/// The list structure enables the analysis to track how loans flow through
-/// different levels of indirection when assignments and dereferences occur.
+///   - For `struct S { View a; Inner b; }` (with `struct Inner { View c; }`),
+///     the node fans out into a tree, one `field` edge per field with origins:
+///       O_s: the record `s` (holds no loans directly)
+///        +-field a-> O_a: what the `View` field `s.a` refers to
+///        +-field b-> O_b: the record `s.b` (holds no loans directly)
+///           +-field c-> O_c: what the `View` field `s.b.c` refers to
 ///
-/// TODO: Currently list-shaped (each node has at most one pointee child).
-/// Will become tree-shaped once field children are added to support
-/// origin trees for records whose fields have origins.
+/// The structure enables the analysis to track how loans flow through
+/// levels of indirection and across record fields when assignments and
+/// dereferences occur.
 class OriginNode {
 public:
+  /// A labeled edge from this node to a child. The label distinguishes how
+  /// the child is reached: a null `FD` means a pointee edge (one level of
+  /// pointer/reference indirection); a non-null `FD` means a field edge
+  /// (the named field of a record). Putting the label on the edge lets
+  /// one child node play different roles per parent. For example, the subtree
+  /// for `s`'s `v` field is reached from `s`'s record (FD=v) and from
+  /// the lvalue outer built for the MemberExpr `s.v` (FD=null).
+  struct Edge {
+    const FieldDecl *FD;
+    OriginNode *Child;
+  };
+
   OriginNode(OriginID OID) : OID(OID) {}
 
+  OriginID getOriginID() const { return OID; }
+
+  llvm::ArrayRef<Edge> children() const { return Children; }
+
+  template <typename Fn> void forEachOrigin(Fn F) const {
+    llvm::SmallVector<const OriginNode *, 4> Worklist{this};
+    while (!Worklist.empty()) {
+      const OriginNode *N = Worklist.pop_back_val();
+      F(N);
+      for (const Edge &E : N->children())
+        Worklist.push_back(E.Child);
+    }
+  }
+
   OriginNode *getPointeeChild() const {
-    return Children.empty() ? nullptr : Children[0];
+    for (const Edge &E : Children)
+      if (!E.FD)
+        return E.Child;
+    return nullptr;
   }
 
-  OriginID getOriginID() const { return OID; }
+  OriginNode *getFieldChild(const FieldDecl *F) const {
+    assert(F);
+    for (const Edge &E : Children)
+      if (E.FD == F)
+        return E.Child;
+    return nullptr;
+  }
 
-  void setChildren(llvm::ArrayRef<OriginNode *> NewChildren) {
+  void setChildren(llvm::ArrayRef<Edge> NewChildren) {
     assert(Children.empty() && "children must be set at most once");
     Children = NewChildren;
   }
----------------
usx95 wrote:

This was confusing. So this is now relevant only for Pointee and reference indirection.

Maybe rename to `getPointeeChainLength`?

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