changeset:   87306:8a2861aaef82
branch:      pep-3154
user:        Alexandre Vassalotti <alexandre@peadrop.com>
date:        Mon Nov 18 22:06:12 2013 -0800
files:       Lib/pickle.py Lib/pickletools.py Modules/_pickle.c
description:
Implement MEMOIZE opcode.


diff -r 1fef6f404e8d -r 8a2861aaef82 Lib/pickle.py
--- a/Lib/pickle.py	Sun Nov 17 20:29:13 2013 -0800
+++ b/Lib/pickle.py	Mon Nov 18 22:06:12 2013 -0800
@@ -174,6 +174,7 @@
 FROZENSET        = b'\x91'  # build frozenset from topmost stack items
 NEWOBJ_EX        = b'\x92'  # like NEWOBJ but work with keyword only arguments
 STACK_GLOBAL     = b'\x93'  # same as GLOBAL but using names on the stacks
+MEMOIZE          = b'\x94'  # store top of the stack in memo
 
 __all__.extend([x for x in dir() if re.match("[A-Z][A-Z0-9_]+$", x)])
 
@@ -441,19 +442,21 @@
         if self.fast:
             return
         assert id(obj) not in self.memo
-        memo_len = len(self.memo)
-        self.write(self.put(memo_len))
-        self.memo[id(obj)] = memo_len, obj
+        idx = len(self.memo)
+        self.write(self.put(idx))
+        self.memo[id(obj)] = idx, obj
 
     # Return a PUT (BINPUT, LONG_BINPUT) opcode string, with argument i.
-    def put(self, i):
-        if self.bin:
-            if i < 256:
-                return BINPUT + pack("<B", i)
+    def put(self, idx):
+        if self.proto >= 4:
+            return MEMOIZE
+        elif self.bin:
+            if idx < 256:
+                return BINPUT + pack("<B", idx)
             else:
-                return LONG_BINPUT + pack("<I", i)
-
-        return PUT + repr(i).encode("ascii") + b'\n'
+                return LONG_BINPUT + pack("<I", idx)
+        else:
+            return PUT + repr(idx).encode("ascii") + b'\n'
 
     # Return a GET (BINGET, LONG_BINGET) opcode string, with argument i.
     def get(self, i):
@@ -1432,6 +1435,11 @@
         self.memo[i] = self.stack[-1]
     dispatch[LONG_BINPUT[0]] = load_long_binput
 
+    def load_memoize(self):
+        memo = self.memo
+        memo[len(memo)] = self.stack[-1]
+    dispatch[MEMOIZE[0]] = load_memoize
+
     def load_append(self):
         stack = self.stack
         value = stack.pop()
diff -r 1fef6f404e8d -r 8a2861aaef82 Lib/pickletools.py
--- a/Lib/pickletools.py	Sun Nov 17 20:29:13 2013 -0800
+++ b/Lib/pickletools.py	Mon Nov 18 22:06:12 2013 -0800
@@ -1816,6 +1816,18 @@
       unsigned little-endian integer following.
       """),
 
+    I(name='MEMOIZE',
+      code='\x94',
+      arg=None,
+      stack_before=[anyobject],
+      stack_after=[anyobject],
+      proto=4,
+      doc="""Store the stack top into the memo.  The stack is not popped.
+
+      The index of the memo location to write is the number of
+      elements currently present in the memo.
+      """),
+
     # Access the extension registry (predefined objects).  Akin to the GET
     # family.
 
@@ -2392,17 +2404,20 @@
                 errormsg = markmsg = "no MARK exists on stack"
 
         # Check for correct memo usage.
-        if opcode.name in ("PUT", "BINPUT", "LONG_BINPUT"):
-            assert arg is not None
-            if arg in memo:
+        if opcode.name in ("PUT", "BINPUT", "LONG_BINPUT", "MEMOIZE"):
+            if opcode.name == "MEMOIZE":
+                memo_idx = len(memo)
+            else:
+                assert arg is not None
+                memo_idx = arg
+            if memo_idx in memo:
                 errormsg = "memo key %r already defined" % arg
             elif not stack:
                 errormsg = "stack is empty -- can't store into memo"
             elif stack[-1] is markobject:
                 errormsg = "can't store markobject in the memo"
             else:
-                memo[arg] = stack[-1]
-
+                memo[memo_idx] = stack[-1]
         elif opcode.name in ("GET", "BINGET", "LONG_BINGET"):
             if arg in memo:
                 assert len(after) == 1
diff -r 1fef6f404e8d -r 8a2861aaef82 Modules/_pickle.c
--- a/Modules/_pickle.c	Sun Nov 17 20:29:13 2013 -0800
+++ b/Modules/_pickle.c	Mon Nov 18 22:06:12 2013 -0800
@@ -81,7 +81,8 @@
     ADDITEMS         = '\x90',
     FROZENSET        = '\x91',
     NEWOBJ_EX        = '\x92',
-    STACK_GLOBAL     = '\x93'
+    STACK_GLOBAL     = '\x93',
+    MEMOIZE          = '\x94'
 };
 
 /* These aren't opcodes -- they're ways to pickle bools before protocol 2
@@ -366,7 +367,8 @@
     /* The unpickler memo is just an array of PyObject *s. Using a dict
        is unnecessary, since the keys are contiguous ints. */
     PyObject **memo;
-    Py_ssize_t memo_size;
+    Py_ssize_t memo_size;       /* Capacity of the memo array */
+    Py_ssize_t memo_len;        /* Number of objects in the memo */
 
     PyObject *arg;
     PyObject *pers_func;        /* persistent_load() method, can be NULL. */
@@ -1261,7 +1263,12 @@
     Py_INCREF(value);
     old_item = self->memo[idx];
     self->memo[idx] = value;
-    Py_XDECREF(old_item);
+    if (old_item != NULL) {
+        Py_DECREF(old_item);
+    }
+    else {
+        self->memo_len++;
+    }
     return 0;
 }
 
@@ -1324,6 +1331,7 @@
     self->fix_imports = 0;
     memset(&self->buffer, 0, sizeof(Py_buffer));
     self->memo_size = 32;
+    self->memo_len = 0;
     self->memo = _Unpickler_NewMemo(self->memo_size);
     self->stack = (Pdata *)Pdata_New();
 
@@ -1438,38 +1446,44 @@
 static int
 memo_put(PicklerObject *self, PyObject *obj)
 {
-    Py_ssize_t x;
     char pdata[30];
     Py_ssize_t len;
-    int status = 0;
+    Py_ssize_t idx;
+
+    const char memoize_op = MEMOIZE;
 
     if (self->fast)
         return 0;
     if (_Pickler_OpcodeBoundary(self))
-        goto error;
-
-    x = PyMemoTable_Size(self->memo);
-    if (PyMemoTable_Set(self->memo, obj, x) < 0)
-        goto error;
-
-    if (!self->bin) {
+        return -1;
+
+    idx = PyMemoTable_Size(self->memo);
+    if (PyMemoTable_Set(self->memo, obj, idx) < 0)
+        return -1;
+
+    if (self->proto >= 4) {
+        if (_Pickler_Write(self, &memoize_op, 1) < 0)
+            return -1;
+        return 0;
+    }
+    else if (!self->bin) {
         pdata[0] = PUT;
         PyOS_snprintf(pdata + 1, sizeof(pdata) - 1,
-                      "%" PY_FORMAT_SIZE_T "d\n", x);
+                      "%" PY_FORMAT_SIZE_T "d\n", idx);
         len = strlen(pdata);
     }
     else {
-        if (x < 256) {
+        if (idx < 256) {
             pdata[0] = BINPUT;
-            pdata[1] = (unsigned char)x;
+            pdata[1] = (unsigned char)idx;
             len = 2;
         }
-        else if (x <= 0xffffffffL) {
+        else if (idx <= 0xffffffffL) {
             pdata[0] = LONG_BINPUT;
-            pdata[1] = (unsigned char)(x & 0xff);
-            pdata[2] = (unsigned char)((x >> 8) & 0xff);
-            pdata[3] = (unsigned char)((x >> 16) & 0xff);
-            pdata[4] = (unsigned char)((x >> 24) & 0xff);
+            pdata[1] = (unsigned char)(idx & 0xff);
+            pdata[2] = (unsigned char)((idx >> 8) & 0xff);
+            pdata[3] = (unsigned char)((idx >> 16) & 0xff);
+            pdata[4] = (unsigned char)((idx >> 24) & 0xff);
             len = 5;
         }
         else { /* unlikely */
@@ -1478,16 +1492,10 @@
             return -1;
         }
     }
-
     if (_Pickler_Write(self, pdata, len) < 0)
-        goto error;
-
-    if (0) {
-  error:
-        status = -1;
-    }
-
-    return status;
+        return -1;
+
+    return 0;
 }
 
 static PyObject *
@@ -5566,6 +5574,18 @@
 }
 
 static int
+load_memoize(UnpicklerObject *self)
+{
+    PyObject *value;
+
+    if (Py_SIZE(self->stack) <= 0)
+        return stack_underflow();
+    value = self->stack->data[Py_SIZE(self->stack) - 1];
+
+    return _Unpickler_MemoPut(self, self->memo_len, value);
+}
+
+static int
 do_append(UnpicklerObject *self, Py_ssize_t x)
 {
     PyObject *value;
@@ -6009,6 +6029,7 @@
         OP(BINPUT, load_binput)
         OP(LONG_BINPUT, load_long_binput)
         OP(PUT, load_put)
+        OP(MEMOIZE, load_memoize)
         OP(POP, load_pop)
         OP(POP_MARK, load_pop_mark)
         OP(SETITEM, load_setitem)