From 8352c3c9bebb5ebc72d1ec7efafbf0e352ab6db6 Mon Sep 17 00:00:00 2001
From: Gleb Koval <gleb@koval.net>
Date: Wed, 10 Jan 2024 12:08:38 +0000
Subject: [PATCH] Implement branches + tests

---
 src/main/kotlin/tinyvm/Repository.kt     | 51 +++++++++++++++++++++---
 src/test/kotlin/tinyvm/RepositoryTest.kt | 29 ++++++++++++++
 2 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/src/main/kotlin/tinyvm/Repository.kt b/src/main/kotlin/tinyvm/Repository.kt
index f48dcc2..a287e31 100644
--- a/src/main/kotlin/tinyvm/Repository.kt
+++ b/src/main/kotlin/tinyvm/Repository.kt
@@ -1,5 +1,7 @@
 package tinyvm
 
+import java.util.TreeSet
+
 class HashCollisionException(hash: String) : Exception("Different object types with identical hash '$hash'")
 
 private class CommitTimeComparator : Comparator<Commit> {
@@ -9,14 +11,49 @@ private class CommitTimeComparator : Comparator<Commit> {
     ): Int = (o1.timestamp.epochSecond - o2.timestamp.epochSecond).toInt()
 }
 
-class Repository {
-    private val commits = sortedSetOf(CommitTimeComparator())
+private val commitTimeComparator = CommitTimeComparator()
+
+private class Branch(val name: String, val commits: TreeSet<Commit>)
+
+class Repository(initialBranch: String = "master") {
+    private var head = Branch(initialBranch, sortedSetOf(commitTimeComparator))
+    private val commits = mutableMapOf(initialBranch to head)
+
+    val branch: String
+        get() = head.name
+
+    val branches: Set<String>
+        get() = commits.keys
 
     // Store all objects in one map like git does. This would simplify the data persistence implementation (if there was
     // one) and allows for other objects to be added in the future without modifying the data persistence implementation
     // at all.
     private val objects = mutableMapOf<String, Object>()
 
+    /**
+     * Create a new branch as a copy of the current branch.
+     * @param name New branch name.
+     * @param use Whether to switch to this branch.
+     */
+    fun createBranch(
+        name: String,
+        use: Boolean = true,
+    ) {
+        commits[name] =
+            Branch(
+                name,
+                head.commits.clone() as? TreeSet<Commit> ?: throw Exception("TreeSet.clone() unexpected result"),
+            )
+        if (use) useBranch(name)
+    }
+
+    /**
+     * Use (switch to) a branch.
+     */
+    fun useBranch(name: String) {
+        head = commits[name] ?: throw BranchNotFoundException(name)
+    }
+
     /**
      * (Deep) get or put a commit object into the repository.
      * This will also get or put all child trees and blobs.
@@ -33,7 +70,7 @@ class Repository {
                 timestamp = commit.timestamp,
             )
         objects[hash] = newCommit
-        commits.add(newCommit)
+        head.commits.add(newCommit)
         return newCommit
     }
 
@@ -45,12 +82,12 @@ class Repository {
     /**
      * List all commits.
      */
-    fun listCommits(): List<Commit> = commits.toList()
+    fun listCommits(): List<Commit> = head.commits.toList()
 
     /**
      * Find commit.
      */
-    fun findCommit(predicate: (Commit) -> Boolean): Commit? = commits.find(predicate)
+    fun findCommit(predicate: (Commit) -> Boolean): Commit? = head.commits.find(predicate)
 
     /**
      * Dump repository objects
@@ -89,4 +126,6 @@ class Repository {
      */
     private inline fun <reified T : Object> findObject(hash: String): T? =
         objects[hash]?.let { it as? T ?: throw HashCollisionException(hash) }
-}
\ No newline at end of file
+}
+
+class BranchNotFoundException(name: String) : Exception("Branch '$name' does not exist")
\ No newline at end of file
diff --git a/src/test/kotlin/tinyvm/RepositoryTest.kt b/src/test/kotlin/tinyvm/RepositoryTest.kt
index 93cb2eb..4a7a7a0 100644
--- a/src/test/kotlin/tinyvm/RepositoryTest.kt
+++ b/src/test/kotlin/tinyvm/RepositoryTest.kt
@@ -118,6 +118,35 @@ internal class RepositoryTest {
         assertEquals(null, repository.getCommit("00000000000000000000"))
     }
 
+    @Test
+    fun `can create and use branch`() {
+        repository.commit(commits[0])
+        repository.commit(commits[1])
+        repository.commit(commits[2])
+        repository.createBranch("test-branch", use = false)
+        assertEquals("master", repository.branch)
+        repository.createBranch("test-branch-2")
+        assertEquals("test-branch-2", repository.branch)
+        assertEquals(setOf("master", "test-branch", "test-branch-2"), repository.branches)
+        assertEquals(commits.map { it.hash() }, repository.listCommits().map { it.hash() })
+        assertEquals(10, repository.dumpObjects().size)
+    }
+
+    @Test
+    fun `can commit to branch`() {
+        repository.commit(commits[0])
+        repository.createBranch("test-branch")
+        repository.commit(commits[1])
+        repository.createBranch("test-branch-2")
+        repository.commit(commits[2])
+
+        assertEquals(commits.map { it.hash() }, repository.listCommits().map { it.hash() })
+        repository.useBranch("test-branch")
+        assertEquals(commits.take(2).map { it.hash() }, repository.listCommits().map { it.hash() })
+        repository.useBranch("master")
+        assertEquals(commits.take(1).map { it.hash() }, repository.listCommits().map { it.hash() })
+    }
+
     @Test
     fun `can display commit`() {
         assertEquals(