Allow all Function objects to implicitly depend on a '__module_hash__' member of their globals.

In some downstream use cases, we end up with Function objects that have been created with
globals dicts that are not backed by actual modules. In this case, if those functions
reference global variables that are mutable, we have no way of differentiating two
functions with identical behavior but different globals.  This is because normally
we don't need to differentiate between them because they get access to their globals
by virtue of their reference to the global module dict, and only one such object
can exist with a given name in memory at once.  But if a function's globals are
not named, we need a way of differentiating the dicts.

To support this, we introduce a __module_hash__ variable - all TP function objects
(Entrypoint, NotCompiled, Function, and Class methods) have an implicit dependency on
this global variable.  This allows downstream code to set this value and prevent
the two functions from being the same.

Author: braxtonmckee

typed_python/FunctionType.hpp

@@ -1076,6 +1076,9 @@ class Function : public Type {
             uint8_t* bytes;
             Py_ssize_t bytecount;
 
+            static PyObject* moduleHashName = PyUnicode_FromString("__module_hash__");
+            outSequences.push_back(std::vector<PyObject*>({moduleHashName}));
+
             PyBytes_AsStringAndSize(((PyCodeObject*)code)->co_code, (char**)&bytes, &bytecount);
 
             long opcodeCount = bytecount / 2;
@@ -1172,6 +1175,8 @@ class Function : public Type {
                     extractGlobalAccessesFromCode((PyCodeObject*)o, outAccesses);
                 }
             });
+
+            outAccesses.insert("__module_hash__");
         }
 
         static void extractNamesFromCode(PyCodeObject* code, std::set<PyObject*>& outNames) {
@@ -1183,6 +1188,9 @@ class Function : public Type {
                     extractNamesFromCode((PyCodeObject*)o, outNames);
                 }
             });
+
+            static PyObject* moduleHashName = PyUnicode_FromString("__module_hash__");
+            outNames.insert(moduleHashName);
         }
 
         void setGlobals(PyObject* globals) {

typed_python/SerializationContext.py

@@ -130,6 +130,10 @@ class SerializationContext(Class, Final):
     nameForObjectOverride = Member(OneOf(None, object))
     objectFromNameOverride = Member(OneOf(None, object))
 
+    # these are used to control how we serialize specific types
+    representationForOverride = Member(OneOf(None, object))
+    setInstanceStateOverride = Member(OneOf(None, object))
+
     def __init__(
         self,
         nameToObjectOverride=None,
@@ -548,6 +552,11 @@ def representationFor(self, inst):
         Returns:
             a representation tuple, or None
         '''
+        if self.representationForOverride is not None:
+            representation = self.representationForOverride(inst)
+            if representation is not None:
+                return representation
+
         if type(inst) in (tuple, list, dict, set, str, int, bool, float):
             return None
 
@@ -682,7 +691,7 @@ def representationFor(self, inst):
             else:
                 globalsToUse = {}
 
-                all_names = set(['__builtins__'])
+                all_names = set(['__builtins__', '__module_hash__'])
 
                 def walkCodeObject(code):
                     all_names.update(code.co_names)
@@ -736,6 +745,12 @@ def walkCodeObject(code):
     def setInstanceStateFromRepresentation(
         self, instance, representation=None, itemIt=None, kvPairIt=None, setStateFun=None
     ):
+        if self.setInstanceStateOverride is not None:
+            if self.setInstanceStateOverride(
+                instance, representation, itemIt, kvPairIt, setStateFun
+            ):
+                return
+
         if representation is reconstructTypeFunctionType:
             return
 

typed_python/SpecialModuleNames.hpp

@@ -91,6 +91,10 @@ inline bool isSpecialIgnorableName(const std::string& name) {
         "__rpow__", "__rrshift__", "__rshift__", "__rsub__",
         "__rtruediv__", "__rxor__", "__setattr__", "__setitem__",
         "__str__", "__sub__", "__truediv__", "__xor__",
+        // this is not a real python magic method, but we want to ensure that
+        // if this is present in a function's globals that we use it to make the
+        // function unique.
+        "__module_hash__"
     });
 
     return (

typed_python/internals.py

@@ -484,6 +484,9 @@ def realizedGlobals(self):
                 if name.split(".")[0] in globalNames:
                     res[name] = self.functionGlobals[name]
 
+            if '__module_hash__' in self.functionGlobals:
+                res['__module_hash__'] = self.functionGlobals['__module_hash__']
+
             for varname, cell in self.funcGlobalsInCells.items():
                 res[varname] = cell.cell_contents
 

typed_python/type_identity_test.py

@@ -14,9 +14,11 @@
 
 import threading
 import pytest
+import tempfile
+import os
 
 import typed_python
-from typed_python.test_util import evaluateExprInFreshProcess
+from typed_python.test_util import evaluateExprInFreshProcess, callFunctionInFreshProcess
 from typed_python import (
     UInt64, UInt32,
     ListOf, TupleOf, Tuple, NamedTuple, Dict, OneOf, Forward, identityHash,
@@ -699,3 +701,86 @@ def f(self):
             return 0
 
     print(typeWalkRecord(N))
+
+
+def test_module_hash_magic_value():
+    with tempfile.TemporaryDirectory() as tempDir:
+
+        def makeFun(mh):
+            """Produce a dummy function with __module_hash__ of 'mh'
+
+            Note that we need to produce this function with 'backing code' so that
+            the AST system can find it and serialize it.
+            """
+            globalsDict = {}
+            fname = os.path.join(tempDir, "code_" + mh + ".py")
+
+            pyCode = (
+                f"from typed_python import Function\n"
+                f"__module_hash__ = '{mh}'\n"
+                f"@Function\n"
+                f"def f(x):\n"
+                f"    return x\n"
+            )
+
+            with open(fname, "w") as f:
+                f.write(pyCode)
+
+            exec(compile(pyCode, fname, "exec"), globalsDict)
+            return globalsDict['f']
+
+        def makeFunIH(mh):
+            return identityHash(type(makeFun(mh)))
+
+        # check that the identity hash depends on the module hash
+        assert identityHash(type(makeFun('A'))) == identityHash(type(makeFun('A')))
+        assert identityHash(type(makeFun('A'))) != identityHash(type(makeFun('B')))
+
+        # functions should have this in their globals
+        f = makeFun('A')
+        assert '__module_hash__' in f.overloads[0].functionGlobals
+        assert '__module_hash__' in f.overloads[0].realizedGlobals
+
+        # even if we execute this in another process, we should get the same function back
+        # and it should have a __module_hash__ in its globals. We have to be careful about
+        # this because we need the serializer to understand that __module_hash__ is special
+        # and that the function implicitly references it.
+        fFromOtherProcess = callFunctionInFreshProcess(makeFun, ('C',))
+        assert fFromOtherProcess.overloads[0].functionGlobals['__module_hash__'] == 'C'
+        assert fFromOtherProcess.overloads[0].realizedGlobals['__module_hash__'] == 'C'
+
+        # check that the identity hash we loaded is the same one we would get from a
+        # subprocess reading it.
+        assert (
+            identityHash(type(fFromOtherProcess))
+            == callFunctionInFreshProcess(makeFunIH, ('C',))
+        )
+
+
+def test_module_hash_magic_value_on_untyped_function_preserved_by_serialization():
+    with tempfile.TemporaryDirectory() as tempDir:
+
+        def makeFun(mh):
+            """Produce a dummy function with __module_hash__ of 'mh'
+
+            Note that we need to produce this function with 'backing code' so that
+            the AST system can find it and serialize it.
+            """
+            globalsDict = {}
+            fname = os.path.join(tempDir, "code_" + mh + ".py")
+
+            pyCode = (
+                f"from typed_python import Entrypoint\n"
+                f"__module_hash__ = '{mh}'\n"
+                f"def f(x):\n"
+                f"    return x\n"
+            )
+
+            with open(fname, "w") as f:
+                f.write(pyCode)
+
+            exec(compile(pyCode, fname, "exec"), globalsDict)
+            return globalsDict['f']
+
+        fFromOtherProcess = callFunctionInFreshProcess(makeFun, ('C',))
+        assert fFromOtherProcess.__globals__['__module_hash__'] == 'C'