-mm=own: a class family spans one class hierarchy, not every class with that slot

tslang/lib/TypeScript/OwnershipSignaturePass.cpp

@@ -53,8 +53,9 @@ namespace own_signature
         findImmortalGlobals(module);
         moveIntoGlobals(module);
         rejectBorrowedGlobalInitializers(module);
-        collectClassVTables(module);
+        // the storages first: a vtable's family is that of its class's root (familyRoot)
         collectClassStorages(module);
+        collectClassVTables(module);
         collectInterfaceFamilies(module);
         fieldFunctions.collect(module);
         // after the class storages (`@cpp`) and the families (holdsFunction reads the interface vtables)

tslang/lib/TypeScript/OwnershipSignaturePass.h

@@ -217,9 +217,10 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
 
     // Class names, from their vtables, to split a method's symbol into class and method name.
     llvm::StringSet<> classNames;
-    // (vtable position, method name) -> the functions at that position under that name, in every
-    // class vtable of the module: what a virtual call through that slot may reach.
-    std::map<std::pair<int64_t, std::string>, llvm::SmallVector<mlir::StringAttr, 2>> families;
+    // (vtable position, method name, root class) -> the functions at that position under that name,
+    // in every class vtable of the module whose class has that root (familyRoot): what a virtual
+    // call through that slot may reach. A slot with a member of an unknown root has one family, root "".
+    std::map<std::tuple<int64_t, std::string, std::string>, llvm::SmallVector<mlir::StringAttr, 2>> families;
 
     // Functions whose callers are not all known - a caller outside the module, or a reference to
     // the symbol that is not a call this pass resolves - and why, for the error in their body.
@@ -630,6 +631,12 @@ class OwnershipSignaturePass : public mlir::PassWrapper<OwnershipSignaturePass,
     // override keeps its base's method name, which is what puts both in one family.
     std::string methodName(llvm::StringRef symbol);
 
+    // The class a class's chain of bases ends in (a base is the first field of a class's storage),
+    // or "" when its storage is not known here or the chain is too deep to follow. A virtual call
+    // reaches only overrides in its class's hierarchy: `Array<string>.find` and `Array<int>.find`,
+    // one generic's specializations, are not one family (they disagreed, and both lost their facts).
+    std::string familyRoot(llvm::StringRef className);
+
     llvm::SmallVector<mlir::StringAttr, 2> familyOf(int64_t index, mlir::StringAttr identifier);
 
     // Adds `symbol`'s function to `callees`; false when it is not defined in this module.

tslang/lib/TypeScript/OwnershipSignatureResolve.cpp

@@ -21,15 +21,46 @@ namespace own_signature
             }
         });
 
+        // each slot's members by the root of the vtable's class; a slot one of whose members has an
+        // unknown root is one family, as every slot was before
+        std::map<std::pair<int64_t, std::string>, llvm::SmallVector<std::pair<std::string, mlir::StringAttr>>> slots;
         for (auto &[name, entries] : vtables)
         {
+            auto root = familyRoot(name.drop_back(llvm::StringRef("..vtbl").size()));
             for (auto [position, symbol] : entries)
             {
-                families[{position, methodName(symbol.getValue())}].push_back(symbol);
+                slots[{position, methodName(symbol.getValue())}].push_back({root, symbol});
+            }
+        }
+
+        for (auto &[slot, members] : slots)
+        {
+            auto unknownRoot = llvm::any_of(members, [](auto &member) { return member.first.empty(); });
+            for (auto &[root, symbol] : members)
+            {
+                families[{slot.first, slot.second, unknownRoot ? std::string() : root}].push_back(symbol);
             }
         }
     }
 
+    std::string OwnershipSignaturePass::familyRoot(llvm::StringRef className)
+    {
+        auto storage = classStorages.lookup(className);
+        for (auto depth = 0; storage && depth < 64; ++depth)
+        {
+            auto fields = storage.getFields();
+            auto base = fields.empty() ? mlir_ts::ClassStorageType() : mlir::dyn_cast<mlir_ts::ClassStorageType>(fields.front().type);
+            if (!base)
+            {
+                return storage.getName().getValue().str();
+            }
+
+            storage = base;
+        }
+
+        return std::string();
+    }
+
     mlir_ts::TupleType OwnershipSignaturePass::storageTupleOf(mlir::Type type)
     {
         return own_facts::storageTupleOf(&getContext(), type);
@@ -88,8 +119,18 @@ namespace own_signature
 
     llvm::SmallVector<mlir::StringAttr, 2> OwnershipSignaturePass::familyOf(int64_t index, mlir::StringAttr identifier)
     {
+        // the family of the identifier's class's root, or the slot's one family of an unknown root
         llvm::SmallVector<mlir::StringAttr, 2> members;
-        if (auto found = families.find({index, methodName(identifier.getValue())}); found != families.end())
+        auto name = methodName(identifier.getValue());
+        auto className =
+            name.size() < identifier.getValue().size() ? identifier.getValue().drop_back(name.size() + 1) : llvm::StringRef();
+        auto found = families.find({index, name, familyRoot(className)});
+        if (found == families.end())
+        {
+            found = families.find({index, name, std::string()});
+        }
+
+        if (found != families.end())
         {
             members = found->second;
         }

tslang/test/tester/CMakeLists.txt

@@ -3955,7 +3955,8 @@ foreach(own_test own_fresh_string own_fresh_array own_return_new own_try_local o
                  own_callback_forwarded own_optional_record_read own_param_default_kept own_literal_into_array
                  own_generator_take own_rest_array_borrows own_string_copies
                  own_captured_string_callback own_callback_and_call own_callback_captured
-                 own_callback_shapes own_callback_extended own_callback_interface)
+                 own_callback_shapes own_callback_extended own_callback_interface
+                 own_generic_specializations)
     tslang_add_test(NAME test-jit-own-${own_test} COMMAND test-runner -jit -mm=own "${PROJECT_SOURCE_DIR}/test/tester/own/${own_test}.ts")
     tslang_add_test(NAME test-compile-own-${own_test} COMMAND test-runner -mm=own "${PROJECT_SOURCE_DIR}/test/tester/own/${own_test}.ts")
     # The promise: no counting at all in what own emits.
@@ -3992,7 +3993,8 @@ foreach(keeps_mm rc gc none)
                        own_optional_record_read own_param_default_kept own_literal_into_array
                        own_generator_take own_rest_array_borrows own_string_copies
                        own_captured_string_callback own_callback_and_call own_callback_captured
-                       own_callback_shapes own_callback_extended own_callback_interface)
+                       own_callback_shapes own_callback_extended own_callback_interface
+                       own_generic_specializations)
         tslang_add_test(NAME test-jit-${keeps_mm}-${keeps_test} COMMAND test-runner -jit -mm=${keeps_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/${keeps_test}.ts")
         tslang_add_test(NAME test-compile-${keeps_mm}-${keeps_test} COMMAND test-runner -mm=${keeps_mm} "${PROJECT_SOURCE_DIR}/test/tester/own/${keeps_test}.ts")
     endforeach()

tslang/test/tester/own/own_generic_specializations.ts

@@ -0,0 +1,30 @@
+// -mm=own: two specializations of one generic class are not one class family. `Box<string>.find`
+// returns a borrow of its field (through `pick`'s result) and `Box<int>.find` a number; neither
+// class extends the other, so a call through either reaches only its own. Grouped by vtable slot and
+// method name alone, they "disagreed" and both lost their facts: with Box<int> beside it,
+// `Box<string>.find` was "'this value' borrows a field and cannot be stored, returned or captured".
+function pick<T>(a: T[], f: (v: T) => boolean): T {
+    for (const v of a) {
+        if (f(v)) return v;
+    }
+
+    return a[0];
+}
+
+class Box<T> {
+    constructor(private data: T[]) {}
+
+    find(f: (v: T) => boolean) {
+        return pick(this.data, f);
+    }
+}
+
+function main() {
+    const s = new Box<string>(["a", "bb"]);
+    const n = new Box<int>([1, 2]);
+
+    assert(s.find((v) => v.length == 2) == "bb");
+    assert(n.find((v) => v > 1) == 2);
+
+    print("done.");
+}