Reuse weakref if no callback is specified

Author: qunaibit

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/WeakRefModuleBuiltins.java

@@ -43,6 +43,7 @@
 import static com.oracle.graal.python.nodes.BuiltinNames.J__WEAKREF;
 import static com.oracle.graal.python.nodes.BuiltinNames.T__WEAKREF;
 import static com.oracle.graal.python.nodes.StringLiterals.T_REF;
+import static com.oracle.graal.python.runtime.exception.PythonErrorType.TypeError;
 
 import java.lang.ref.Reference;
 import java.lang.ref.ReferenceQueue;
@@ -70,6 +71,7 @@
 import com.oracle.graal.python.nodes.PGuards;
 import com.oracle.graal.python.nodes.WriteUnraisableNode;
 import com.oracle.graal.python.nodes.attributes.ReadAttributeFromObjectNode;
+import com.oracle.graal.python.nodes.attributes.WriteAttributeToDynamicObjectNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinBaseNode;
 import com.oracle.graal.python.nodes.function.PythonBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonTernaryBuiltinNode;
@@ -79,7 +81,6 @@
 import com.oracle.graal.python.runtime.PythonContext;
 import com.oracle.graal.python.runtime.PythonOptions;
 import com.oracle.graal.python.runtime.exception.PException;
-import com.oracle.graal.python.runtime.exception.PythonErrorType;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.dsl.Bind;
 import com.oracle.truffle.api.dsl.Cached;
@@ -95,11 +96,58 @@ public class WeakRefModuleBuiltins extends PythonBuiltins {
     private static final HiddenKey weakRefQueueKey = new HiddenKey("weakRefQueue");
     private final ReferenceQueue<Object> weakRefQueue = new ReferenceQueue<>();
 
+    // This GraalPy specific as CPython is storing weakref list within a PyObject (obj +
+    // tp_weaklistoffset)
+    public static final HiddenKey __WEAKLIST__ = new HiddenKey("__weaklist__");
+
     @Override
     protected List<? extends NodeFactory<? extends PythonBuiltinBaseNode>> getNodeFactories() {
         return WeakRefModuleBuiltinsFactory.getFactories();
     }
 
+    private static long getBuiltinTypeWeaklistoffset(PythonBuiltinClassType cls) {
+        // @formatter:off
+        return switch (cls) {
+            case PythonClass -> 368; // type
+            case PSet, // set
+                    PFrozenSet // frozenset
+                    -> 192;
+            case PMemoryView, // memoryview
+                    PCode // code
+                    -> 136;
+            case PBuiltinFunctionOrMethod, // builtin_function_or_method
+                    PGenerator, // generator
+                    PCoroutine, // coroutine
+                    PythonModule, // module
+                    PThreadLocal, // _thread._local
+                    PRLock, // _thread.RLock
+                    PBufferedRWPair, // _io.BufferedRWPair
+                    PAsyncGenerator // async_generator
+                    -> 40;
+            case PMethod, // method
+                    PFileIO, // _io.FileIO
+                    PTee // itertools._tee
+                    -> 32;
+            case PFunction -> 80; // function
+            case PIOBase, // _io._IOBase
+                    PRawIOBase, // _io._RawIOBase
+                    PBufferedIOBase, // _io._BufferedIOBase
+                    PTextIOBase // _io._TextIOBase
+                    -> 24;
+            case PBytesIO, // _io.BytesIO
+                    PPartial // functools.partial
+                    -> 48;
+            case PStringIO -> 112; // _io.StringIO
+            case PBufferedReader, // _io.BufferedReader
+                    PBufferedWriter, // _io.BufferedWriter
+                    PBufferedRandom // _io.BufferedRandom
+                    -> 144;
+            case PTextIOWrapper -> 176; // _io.TextIOWrapper
+            default -> 0;
+            // @formatter:on
+        };
+    }
+
     private static class WeakrefCallbackAction extends AsyncHandler.AsyncPythonAction {
         private final WeakRefStorage[] references;
         private int index;
@@ -179,8 +227,31 @@ public abstract static class ReferenceTypeNode extends PythonTernaryBuiltinNode
         @Child private CExtNodes.GetTypeMemberNode getTpWeaklistoffsetNode;
 
         @Specialization(guards = "!isNativeObject(object)")
-        public PReferenceType refType(Object cls, Object object, @SuppressWarnings("unused") PNone none) {
-            return factory().createReferenceType(cls, object, null, getWeakReferenceQueue());
+        public PReferenceType refType(Object cls, Object object, @SuppressWarnings("unused") PNone none,
+                        @Cached InlinedGetClassNode getClassNode,
+                        @Cached ReadAttributeFromObjectNode getAttrNode,
+                        @Cached WriteAttributeToDynamicObjectNode setAttrNode) {
+            Object obj = object;
+            if (object instanceof PythonBuiltinClassType tobj) {
+                obj = getContext().getCore().lookupType(tobj);
+            }
+
+            Object clazz = getClassNode.execute(this, obj);
+            boolean allowed = true;
+            if (clazz instanceof PythonBuiltinClassType type) {
+                allowed = getBuiltinTypeWeaklistoffset(type) != 0;
+            }
+            if (!allowed) {
+                throw raise(TypeError, ErrorMessages.CANNOT_CREATE_WEAK_REFERENCE_TO, obj);
+            }
+            Object wr = getAttrNode.execute(obj, __WEAKLIST__);
+            if (wr != PNone.NO_VALUE) {
+                return (PReferenceType) wr; // is must be a PReferenceType instance.
+            }
+
+            PReferenceType ref = factory().createReferenceType(cls, obj, null, getWeakReferenceQueue());
+            setAttrNode.execute(obj, __WEAKLIST__, ref);
+            return ref;
         }
 
         @Specialization(guards = {"!isNativeObject(object)", "!isPNone(callback)"})
@@ -223,7 +294,7 @@ public PReferenceType refType(Object cls, PythonAbstractNativeObject pythonObjec
 
         @Fallback
         public PReferenceType refType(@SuppressWarnings("unused") Object cls, Object object, @SuppressWarnings("unused") Object callback) {
-            throw raise(PythonErrorType.TypeError, ErrorMessages.CANNOT_CREATE_WEAK_REFERENCE_TO, object);
+            throw raise(TypeError, ErrorMessages.CANNOT_CREATE_WEAK_REFERENCE_TO, object);
         }
 
         @SuppressWarnings("unchecked")