Allow writing a NoteMap back to the repository

This is necessary to allow applications to wrap the note tree in
a commit and update the note branch with the new state.

Change-Id: Idbd7ead4a1b16ae2b64a30a4a01a29cfed548cdf
Signed-off-by: Shawn O. Pearce <spearce@spearce.org>

5 files changed, 131 insertions, 0 deletions

diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/notes/FanoutBucket.java b/org.eclipse.jgit/src/org/eclipse/jgit/notes/FanoutBucket.java
index 760528618c..2085676900 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/notes/FanoutBucket.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/notes/FanoutBucket.java
@@ -43,12 +43,16 @@
 
 package org.eclipse.jgit.notes;
 
+import static org.eclipse.jgit.lib.FileMode.TREE;
+
 import java.io.IOException;
 
 import org.eclipse.jgit.lib.AbbreviatedObjectId;
 import org.eclipse.jgit.lib.AnyObjectId;
 import org.eclipse.jgit.lib.ObjectId;
+import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.ObjectReader;
+import org.eclipse.jgit.lib.TreeFormatter;
 
 /**
  * A note tree holding only note subtrees, each named using a 2 digit hex name.
@@ -136,6 +140,43 @@ class FanoutBucket extends InMemoryNoteBucket {
 		}
 	}
 
+	private static final byte[] hexchar = { '0', '1', '2', '3', '4', '5', '6',
+			'7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' };
+
+	@Override
+	ObjectId writeTree(ObjectInserter inserter) throws IOException {
+		byte[] nameBuf = new byte[2];
+		TreeFormatter fmt = new TreeFormatter(treeSize());
+		NonNoteEntry e = nonNotes;
+
+		for (int cell = 0; cell < 256; cell++) {
+			NoteBucket b = table[cell];
+			if (b == null)
+				continue;
+
+			nameBuf[0] = hexchar[cell >>> 4];
+			nameBuf[1] = hexchar[cell & 0x0f];
+
+			while (e != null && e.pathCompare(nameBuf, 0, 2, TREE) < 0) {
+				e.format(fmt);
+				e = e.next;
+			}
+
+			fmt.append(nameBuf, 0, 2, TREE, b.writeTree(inserter));
+		}
+
+		for (; e != null; e = e.next)
+			e.format(fmt);
+		return fmt.insert(inserter);
+	}
+
+	private int treeSize() {
+		int sz = cnt * TreeFormatter.entrySize(TREE, 2);
+		for (NonNoteEntry e = nonNotes; e != null; e = e.next)
+			sz += e.treeEntrySize();
+		return sz;
+	}
+
 	private int cell(AnyObjectId id) {
 		return id.getByte(prefixLen >> 1);
 	}
@@ -158,6 +199,11 @@ class FanoutBucket extends InMemoryNoteBucket {
 			return load(noteOn, or).set(noteOn, noteData, or);
 		}
 
+		@Override
+		ObjectId writeTree(ObjectInserter inserter) {
+			return treeId;
+		}
+
 		private NoteBucket load(AnyObjectId objId, ObjectReader or)
 				throws IOException {
 			AbbreviatedObjectId p = objId.abbreviate(prefixLen + 2);
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/notes/LeafBucket.java b/org.eclipse.jgit/src/org/eclipse/jgit/notes/LeafBucket.java
index ce4feae259..40be45f433 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/notes/LeafBucket.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/notes/LeafBucket.java
@@ -43,11 +43,16 @@
 
 package org.eclipse.jgit.notes;
 
+import static org.eclipse.jgit.lib.Constants.OBJECT_ID_STRING_LENGTH;
+import static org.eclipse.jgit.lib.FileMode.REGULAR_FILE;
+
 import java.io.IOException;
 
 import org.eclipse.jgit.lib.AnyObjectId;
 import org.eclipse.jgit.lib.ObjectId;
+import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.ObjectReader;
+import org.eclipse.jgit.lib.TreeFormatter;
 
 /**
  * A note tree holding only notes, with no subtrees.
@@ -126,6 +131,39 @@ class LeafBucket extends InMemoryNoteBucket {
 		}
 	}
 
+	@Override
+	ObjectId writeTree(ObjectInserter inserter) throws IOException {
+		byte[] nameBuf = new byte[OBJECT_ID_STRING_LENGTH];
+		int nameLen = OBJECT_ID_STRING_LENGTH - prefixLen;
+		TreeFormatter fmt = new TreeFormatter(treeSize(nameLen));
+		NonNoteEntry e = nonNotes;
+
+		for (int i = 0; i < cnt; i++) {
+			Note n = notes[i];
+
+			n.copyTo(nameBuf, 0);
+
+			while (e != null
+					&& e.pathCompare(nameBuf, prefixLen, nameLen, REGULAR_FILE) < 0) {
+				e.format(fmt);
+				e = e.next;
+			}
+
+			fmt.append(nameBuf, prefixLen, nameLen, REGULAR_FILE, n.getData());
+		}
+
+		for (; e != null; e = e.next)
+			e.format(fmt);
+		return fmt.insert(inserter);
+	}
+
+	private int treeSize(final int nameLen) {
+		int sz = cnt * TreeFormatter.entrySize(REGULAR_FILE, nameLen);
+		for (NonNoteEntry e = nonNotes; e != null; e = e.next)
+			sz += e.treeEntrySize();
+		return sz;
+	}
+
 	void parseOneEntry(AnyObjectId noteOn, AnyObjectId noteData) {
 		growIfFull();
 		notes[cnt++] = new Note(noteOn, noteData.copy());
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NonNoteEntry.java b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NonNoteEntry.java
index 20fd3a88ad..6a2d44bca9 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NonNoteEntry.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NonNoteEntry.java
@@ -68,4 +68,33 @@ class NonNoteEntry extends ObjectId {
 	void format(TreeFormatter fmt) {
 		fmt.append(name, mode, this);
 	}
+
+	int treeEntrySize() {
+		return TreeFormatter.entrySize(mode, name.length);
+	}
+
+	int pathCompare(byte[] bBuf, int bPos, int bLen, FileMode bMode) {
+		return pathCompare(name, 0, name.length, mode, //
+				bBuf, bPos, bLen, bMode);
+	}
+
+	private static int pathCompare(final byte[] aBuf, int aPos, final int aEnd,
+			final FileMode aMode, final byte[] bBuf, int bPos, final int bEnd,
+			final FileMode bMode) {
+		while (aPos < aEnd && bPos < bEnd) {
+			int cmp = (aBuf[aPos++] & 0xff) - (bBuf[bPos++] & 0xff);
+			if (cmp != 0)
+				return cmp;
+		}
+
+		if (aPos < aEnd)
+			return (aBuf[aPos] & 0xff) - lastPathChar(bMode);
+		if (bPos < bEnd)
+			return lastPathChar(aMode) - (bBuf[bPos] & 0xff);
+		return 0;
+	}
+
+	private static int lastPathChar(final FileMode mode) {
+		return FileMode.TREE.equals(mode.getBits()) ? '/' : '\0';
+	}
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteBucket.java b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteBucket.java
index f75fc1f942..f5465148a8 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteBucket.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteBucket.java
@@ -47,6 +47,7 @@ import java.io.IOException;
 
 import org.eclipse.jgit.lib.AnyObjectId;
 import org.eclipse.jgit.lib.ObjectId;
+import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.ObjectReader;
 
 /**
@@ -61,4 +62,6 @@ abstract class NoteBucket {
 
 	abstract InMemoryNoteBucket set(AnyObjectId noteOn, AnyObjectId noteData,
 			ObjectReader reader) throws IOException;
+
+	abstract ObjectId writeTree(ObjectInserter inserter) throws IOException;
 }
diff --git a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteMap.java b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteMap.java
index 88ab481794..6a7b5cffbe 100644
--- a/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteMap.java
+++ b/org.eclipse.jgit/src/org/eclipse/jgit/notes/NoteMap.java
@@ -309,6 +309,21 @@ public class NoteMap {
 		set(noteOn, null);
 	}
 
+	/**
+	 * Write this note map as a tree.
+	 *
+	 * @param inserter
+	 *            inserter to use when writing trees to the object database.
+	 *            Caller is responsible for flushing the inserter before trying
+	 *            to read the objects, or exposing them through a reference.
+	 * @return the top level tree.
+	 * @throws IOException
+	 *             a tree could not be written.
+	 */
+	public ObjectId writeTree(ObjectInserter inserter) throws IOException {
+		return root.writeTree(inserter);
+	}
+
 	private void load(ObjectId rootTree) throws MissingObjectException,
 			IncorrectObjectTypeException, CorruptObjectException, IOException {
 		AbbreviatedObjectId none = AbbreviatedObjectId.fromString("");