From d740fa5acfced4066a74668452ca6535507c4d1f Mon Sep 17 00:00:00 2001
From: Yuriy Husnay <husnaj23@gmail.com>
Date: Wed, 4 May 2022 18:04:41 +0300
Subject: [PATCH] add support for branch ordering

---
 src/diagrams/git/gitGraphAst.js           | 25 ++++++++++----
 src/diagrams/git/gitGraphParserV2.spec.js | 40 +++++++++++++++++++++++
 src/diagrams/git/parser/gitGraph.jison    |  8 ++++-
 3 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/src/diagrams/git/gitGraphAst.js b/src/diagrams/git/gitGraphAst.js
index 1beb1a89e7..84f27de0eb 100644
--- a/src/diagrams/git/gitGraphAst.js
+++ b/src/diagrams/git/gitGraphAst.js
@@ -8,6 +8,8 @@ import common from '../common/common';
 let mainBranchName = getConfig().gitGraph.mainBranchName;
 let commits = {};
 let head = null;
+let branchesConfig = {};
+branchesConfig[mainBranchName] = { name: mainBranchName, order: 0 };
 let branches = {};
 branches[mainBranchName] = head;
 let curBranch = mainBranchName;
@@ -113,10 +115,11 @@ export const commit = function (msg, id, type, tag) {
   log.debug('in pushCommit ' + commit.id);
 };
 
-export const branch = function (name) {
+export const branch = function (name, order) {
   name = common.sanitizeText(name, configApi.getConfig());
   if (typeof branches[name] === 'undefined') {
     branches[name] = head != null ? head.id : null;
+    branchesConfig[name] = { name, order: order ? parseInt(order, 10) : null };
     checkout(name);
     log.debug('in createBranch');
   } else {
@@ -324,17 +327,25 @@ export const clear = function () {
   let mainBranch = getConfig().gitGraph.mainBranchName;
   branches = {};
   branches[mainBranch] = null;
+  branchesConfig = {};
+  branchesConfig[mainBranch] = { name: mainBranch, order: 0 };
   curBranch = mainBranch;
   seq = 0;
 };
 
 export const getBranchesAsObjArray = function () {
-  const branchArr = [];
-  for (let branch in branches) {
-    // branchArr.push({ name: branch, commit: commits[branches[branch]] });
-    branchArr.push({ name: branch });
-  }
-  return branchArr;
+  const branchesArray = Object.values(branchesConfig)
+    .map((branchConfig, i) => {
+      if (branchConfig.order !== null) return branchConfig;
+      return {
+        ...branchConfig,
+        order: parseFloat(`0.${i}`, 10),
+      };
+    })
+    .sort((a, b) => a.order - b.order)
+    .map(({ name }) => ({ name }));
+
+  return branchesArray;
 };
 
 export const getBranches = function () {
diff --git a/src/diagrams/git/gitGraphParserV2.spec.js b/src/diagrams/git/gitGraphParserV2.spec.js
index 670c3a371d..978528a7e4 100644
--- a/src/diagrams/git/gitGraphParserV2.spec.js
+++ b/src/diagrams/git/gitGraphParserV2.spec.js
@@ -364,6 +364,46 @@ describe('when parsing a gitGraph', function () {
     expect(parser.yy.getDirection()).toBe('LR');
     expect(Object.keys(parser.yy.getBranches()).length).toBe(2);
   });
+  it('should handle new branch checkout with order', function () {
+    const str = `gitGraph:
+    commit
+    branch test1 order: 3
+    branch test2 order: 2
+    branch test3 order: 1
+    `;
+
+    parser.parse(str);
+    const commits = parser.yy.getCommits();
+    expect(Object.keys(commits).length).toBe(1);
+    expect(parser.yy.getCurrentBranch()).toBe('test3');
+    expect(Object.keys(parser.yy.getBranches()).length).toBe(4);
+    expect(parser.yy.getBranchesAsObjArray()).toStrictEqual([
+      { name: 'main' },
+      { name: 'test3' },
+      { name: 'test2' },
+      { name: 'test1' },
+    ]);
+  });
+  it('should handle new branch checkout with and without order', function () {
+    const str = `gitGraph:
+    commit
+    branch test1 order: 1
+    branch test2
+    branch test3
+    `;
+
+    parser.parse(str);
+    const commits = parser.yy.getCommits();
+    expect(Object.keys(commits).length).toBe(1);
+    expect(parser.yy.getCurrentBranch()).toBe('test3');
+    expect(Object.keys(parser.yy.getBranches()).length).toBe(4);
+    expect(parser.yy.getBranchesAsObjArray()).toStrictEqual([
+      { name: 'main' },
+      { name: 'test2' },
+      { name: 'test3' },
+      { name: 'test1' },
+    ]);
+  });
 
   it('should handle new branch checkout & commit', function () {
     const str = `gitGraph:
diff --git a/src/diagrams/git/parser/gitGraph.jison b/src/diagrams/git/parser/gitGraph.jison
index 7c11051443..4ac2a603f5 100644
--- a/src/diagrams/git/parser/gitGraph.jison
+++ b/src/diagrams/git/parser/gitGraph.jison
@@ -36,6 +36,7 @@
 "HIGHLIGHT"                             return 'HIGHLIGHT';
 "tag:"                                  return 'COMMIT_TAG';
 "branch"                                return 'BRANCH';
+"order:"                                return 'ORDER';
 "merge"                                 return 'MERGE';
 // "reset"                                 return 'RESET';
 "checkout"                              return 'CHECKOUT';
@@ -49,6 +50,7 @@
 ["]                                     this.begin("string");
 <string>["]                             this.popState();
 <string>[^"]*                           return 'STR';
+[0-9]+                                  return 'NUM';
 [a-zA-Z][-_\./a-zA-Z0-9]*[-_a-zA-Z0-9]  return 'ID';
 <<EOF>>                                 return 'EOF';
 
@@ -90,10 +92,14 @@ line
 statement
     : commitStatement
     | mergeStatement
-    | BRANCH ID {yy.branch($2)}
+    | branchStatement
     | CHECKOUT ID {yy.checkout($2)}
     // | RESET reset_arg {yy.reset($2)}
     ;
+branchStatement
+    : BRANCH ID {yy.branch($2)}
+    | BRANCH ID ORDER NUM {yy.branch($2, $4)}
+    ;
 
 mergeStatement
     : MERGE ID {yy.merge($2)}