From b10ba0327d10a0621c3e18249ec3e6959c0c7973 Mon Sep 17 00:00:00 2001
From: Durham Goode <durham@fb.com>
Date: Mon, 7 Aug 2017 19:27:17 -0700
Subject: [PATCH] fastmanifest: move tree writing to memmanifestctx.write()

Summary:
Previously fastmanifest hooked into manifestrevlog.add() to write the local tree
manifest when a flat manifest is written. In a future diff we'll be moving the
datastore objects off of the opener object and on to the manifestlog. Since we
need the datastore to write trees, and the datastore will soon be on the
manifestlog, and the manifestlog isn't accessible from manifestrevlog.add(), we
need to move our tree write up to the ctx level instead. This is probably
cleaner anyway since the ctx level is where we differentiate manifest types, not
at the revlog/storage level.

Test Plan: Ran the tests

Reviewers: #fbhgext, ryanmce

Reviewed By: #fbhgext, ryanmce

Differential Revision: https://phab.mercurial-scm.org/D233
---
 fastmanifest/__init__.py       |  2 ++
 fastmanifest/implementation.py | 35 +++++++++++++++++++++-------------
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/fastmanifest/__init__.py b/fastmanifest/__init__.py
index adbeca3085..658c7b511f 100644
--- a/fastmanifest/__init__.py
+++ b/fastmanifest/__init__.py
@@ -202,6 +202,8 @@ class FastManifestExtension(object):
             manifest.manifestlog, '__getitem__', factory.newgetitem)
         extensions.wrapfunction(
             manifest.manifestlog, 'get', factory.newgetdirmanifestctx)
+        extensions.wrapfunction(
+            manifest.memmanifestctx, 'write', factory.ctxwrite)
         extensions.wrapfunction(manifest.manifestrevlog, 'add', factory.add)
 
         if ui.configbool('fastmanifest', 'usecache', True):
diff --git a/fastmanifest/implementation.py b/fastmanifest/implementation.py
index dcf540de1a..976444a9e2 100644
--- a/fastmanifest/implementation.py
+++ b/fastmanifest/implementation.py
@@ -968,18 +968,27 @@ class manifestfactory(object):
             # If neither cache could help, fallback to the normal add
             node = orig(*args, **kwargs)
 
-        # Add the new manifest to the tree
+        return node
+
+    def ctxwrite(self, orig, mfctx, transaction, link, p1, p2, added, removed):
+        node = orig(mfctx, transaction, link, p1, p2, added, removed)
+
+        treeenabled = self.ui.configbool("fastmanifest", "usetree", False)
         if supportsctree and treeenabled:
+            mfl = mfctx._manifestlog
+            datastore = mfctx._revlog().opener.manifestdatastore
+            opener = mfctx._revlog().opener
+
+            m = mfctx._manifestdict
+
+            # hybridmanifest requires you provide either a value or loadflat.
+            # Let's give it a dummy value, since we know we'll only be calling
+            # _treemanifest()
             def loadflat():
-                # This should eventually be made lazy loaded, so consumers can
-                # access the node/p1/linkrev data without having to parse the
-                # whole manifest.
-                data = origself.revision(p1)
-                arraytext = bytearray(data)
-                origself._fulltextcache[p1] = arraytext
-                return manifest.manifestdict(data)
-            tree = hybridmanifest(self.ui, origself.opener, loadflat=loadflat,
+                raise RuntimeError("no-op loadflat should never be hit")
+            tree = hybridmanifest(self.ui, opener, mfl, loadflat=loadflat,
                                   node=p1)._treemanifest()
+
             if tree is not None:
                 newtree = tree.copy()
                 for filename in removed:
@@ -992,7 +1001,7 @@ class manifestfactory(object):
 
                 if not util.safehasattr(transaction, 'treedatapack'):
                     packpath = shallowutil.getlocalpackpath(
-                            origself.opener.vfs.base,
+                            opener.vfs.base,
                             'manifests')
                     transaction.treedatapack = datapack.mutabledatapack(
                             self.ui,
@@ -1003,9 +1012,9 @@ class manifestfactory(object):
                     def finalize(tr):
                         tr.treedatapack.close()
                         tr.treehistpack.close()
-                        treemanifestcache.getinstance(origself.opener,
+                        treemanifestcache.getinstance(opener,
                                                       self.ui).clear()
-                        origself.opener.manifestdatastore.markforrefresh()
+                        datastore.markforrefresh()
                     def abort(tr):
                         tr.treedatapack.abort()
                         tr.treehistpack.abort()
@@ -1024,7 +1033,7 @@ class manifestfactory(object):
                     dpack.add(nname, nnode, revlog.nullid, ntext)
                     hpack.add(nname, nnode, np1, np2, revlog.nullid, '')
 
-                treemanifestcache.getinstance(origself.opener,
+                treemanifestcache.getinstance(opener,
                                               self.ui)[node] = newtree
 
         return node