hg: allow writing to inmemory packs on bundlerepo

Summary:
In a future diff we'll want to allow converting flat manifests to tree
manifests during infinitepush. To do so, let's allow writing to in memory packs
on bundlerepositories.

Reviewed By: StanislavGlebik

Differential Revision: D7256563

fbshipit-source-id: 10ec58d1171b7882d6db9a916c50a19bc11dbcb4

hgext/remotefilelog/datapack.py

@@ -598,3 +598,30 @@ class mutabledatapack(basepack.mutablebasepack):
                 missing.append((name, node))
 
         return missing
+
+class memdatapack(object):
+    def __init__(self):
+        self.data = {}
+        self.meta = {}
+
+    def add(self, name, node, deltabase, delta):
+        self.data[(name, node)] = (deltabase, delta)
+
+    def getdelta(self, name, node):
+        deltabase, delta = self.data[(name, node)]
+        return (delta, name, deltabase,
+                self.getmeta(name, node))
+
+    def getdeltachain(self, name, node):
+        deltabase, delta = self.data[(name, node)]
+        return [(name, node, name, deltabase, delta)]
+
+    def getmeta(self, name, node):
+        return self.meta[(name, node)]
+
+    def getmissing(self, keys):
+        missing = []
+        for key in keys:
+            if key not in self.data:
+                missing.append(key)
+        return missing

hgext/remotefilelog/historypack.py

@@ -565,3 +565,35 @@ class mutablehistorypack(basepack.mutablebasepack):
                 missing.append((name, node))
 
         return missing
+
+class memhistorypack(object):
+    def __init__(self):
+        self.history = {}
+
+    def add(self, name, node, p1, p2, linknode, copyfrom):
+        self.history.setdefault(name, {})[node] = (p1, p2, linknode, copyfrom)
+
+    def getmissing(self, keys):
+        missing = []
+        for name, node in keys:
+            filehistory = self.history.get(name)
+            if filehistory is None:
+                missing.append((name, node))
+            else:
+                if node not in filehistory:
+                    missing.append((name, node))
+        return missing
+
+    def getancestors(self, name, node, known=None):
+        ancestors = {}
+        try:
+            ancestors[node] = self.history[name][node]
+        except KeyError:
+            raise KeyError((name, node))
+        return ancestors
+
+    def getnodeinfo(self, name, node):
+        try:
+            return self.history[name][node]
+        except KeyError:
+            raise KeyError((name, node))

hgext/treemanifest/__init__.py

@@ -141,11 +141,13 @@ from ..remotefilelog.metadatastore import (
 from ..remotefilelog.datapack import (
     datapack,
     datapackstore,
+    memdatapack,
     mutabledatapack,
 )
 from ..remotefilelog.historypack import (
     historypack,
     historypackstore,
+    memhistorypack,
     mutablehistorypack,
 )
 from ..remotefilelog.repack import (
@@ -199,6 +201,8 @@ def uisetup(ui):
     extensions.wrapfunction(revlog.revlog, 'checkhash', _checkhash)
 
     wrappropertycache(localrepo.localrepository, 'manifestlog', getmanifestlog)
+    wrappropertycache(bundlerepo.bundlerepository, 'manifestlog',
+                      getbundlemanifestlog)
 
     extensions.wrapfunction(
         manifest.memmanifestctx, 'write', _writemanifestwrapper)
@@ -522,14 +526,17 @@ class basetreemanifestlog(object):
             return self._addtopack(ui, newtree, p1node, p2node,
                     overridenode=overridenode, overridep1node=overridep1node)
 
+    def _createmutablepack(self):
+        packpath = shallowutil.getlocalpackpath(
+                self._opener.vfs.base,
+                'manifests')
+        self._mutabledatapack = mutabledatapack(self.ui, packpath)
+        self._mutablehistorypack = mutablehistorypack(self.ui, packpath)
+
     def _addtopack(self, ui, newtree, p1node, p2node, overridenode=None,
                    overridep1node=None):
         if self._mutabledatapack is None:
-            packpath = shallowutil.getlocalpackpath(
-                    self._opener.vfs.base,
-                    'manifests')
-            self._mutabledatapack = mutabledatapack(ui, packpath)
-            self._mutablehistorypack = mutablehistorypack(ui, packpath)
+            self._createmutablepack()
 
         p1tree = self[p1node].read()
         if util.safehasattr(p1tree, '_treemanifest'):
@@ -854,6 +861,35 @@ def getmanifestlog(orig, self):
 
     return mfl
 
+def getbundlemanifestlog(orig, self):
+    mfl = orig(self)
+    if not treeenabled(self.ui):
+        return mfl
+
+    wrapmfl = mfl
+    if isinstance(mfl, hybridmanifestlog):
+        wrapmfl = mfl.treemanifestlog
+
+    class bundlemanifestlog(wrapmfl.__class__):
+        def add(self, ui, newtree, p1node, p2node, overridenode=None,
+                overridep1node=None, tr=None, linkrev=None):
+            return self._addtopack(ui, newtree, p1node, p2node,
+                    overridenode=overridenode, overridep1node=overridep1node)
+
+        def _createmutablepack(self):
+            self._mutabledatapack = memdatapack()
+            self._mutablehistorypack = memhistorypack()
+
+        def commitpending(self):
+            pass
+
+        def abortpending(self):
+            self._mutabledatapack = None
+            self._mutablehistorypack = None
+
+    wrapmfl.__class__ = bundlemanifestlog
+    return mfl
+
 def _writemanifestwrapper(orig, self, tr, link, p1, p2, added, removed):
     n = orig(self, tr, link, p1, p2, added, removed)
 

mercurial/bundlerepo.py

@@ -370,6 +370,10 @@ class bundlerepository(localrepo.localrepository):
         self.manstart = self._cgunpacker.tell()
         return c
 
+    @localrepo.unfilteredpropertycache
+    def manifestlog(self):
+        return super(bundlerepository, self).manifestlog
+
     def _constructmanifest(self):
         self._cgunpacker.seek(self.manstart)
         # consume the header if it exists