From 328601be782bb0b43dd7f6029d6a3ea3f9c54a7a Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Tue, 8 Aug 2023 09:09:05 +0200
Subject: [PATCH] fix(template): save Segments correctly

resolves #4143
---
 src/platform/cache.go                      |  2 +-
 src/platform/{concurrent_map.go => map.go} | 12 +++++++-
 src/platform/shell.go                      | 36 ++++++++++++----------
 src/template/text.go                       |  2 +-
 src/template/text_test.go                  |  2 +-
 5 files changed, 34 insertions(+), 20 deletions(-)
 rename src/platform/{concurrent_map.go => map.go} (75%)

diff --git a/src/platform/cache.go b/src/platform/cache.go
index c80ea4f7..0e8298ba 100644
--- a/src/platform/cache.go
+++ b/src/platform/cache.go
@@ -56,7 +56,7 @@ func (fc *fileCache) Close() {
 	if !fc.dirty {
 		return
 	}
-	cache := fc.cache.List()
+	cache := fc.cache.SimpleMap()
 	if dump, err := json.MarshalIndent(cache, "", "    "); err == nil {
 		cacheFilePath := filepath.Join(fc.cachePath, CacheFile)
 		_ = os.WriteFile(cacheFilePath, dump, 0644)
diff --git a/src/platform/concurrent_map.go b/src/platform/map.go
similarity index 75%
rename from src/platform/concurrent_map.go
rename to src/platform/map.go
index 5b4a780c..fe929d34 100644
--- a/src/platform/concurrent_map.go
+++ b/src/platform/map.go
@@ -26,7 +26,7 @@ func (cm *ConcurrentMap) Contains(key string) bool {
 	return ok
 }
 
-func (cm *ConcurrentMap) List() map[string]any {
+func (cm *ConcurrentMap) SimpleMap() SimpleMap {
 	list := make(map[string]any)
 	(*sync.Map)(cm).Range(func(key, value any) bool {
 		list[key.(string)] = value
@@ -34,3 +34,13 @@ func (cm *ConcurrentMap) List() map[string]any {
 	})
 	return list
 }
+
+type SimpleMap map[string]any
+
+func (m SimpleMap) ConcurrentMap() *ConcurrentMap {
+	var cm ConcurrentMap
+	for k, v := range m {
+		cm.Set(k, v)
+	}
+	return &cm
+}
diff --git a/src/platform/shell.go b/src/platform/shell.go
index dd8dffb9..c8c6a3bf 100644
--- a/src/platform/shell.go
+++ b/src/platform/shell.go
@@ -168,21 +168,22 @@ type SystemInfo struct {
 }
 
 type TemplateCache struct {
-	Root         bool
-	PWD          string
-	Folder       string
-	Shell        string
-	ShellVersion string
-	UserName     string
-	HostName     string
-	Code         int
-	Env          map[string]string
-	Var          map[string]interface{}
-	OS           string
-	WSL          bool
-	PromptCount  int
-	SHLVL        int
-	Segments     *ConcurrentMap
+	Root          bool
+	PWD           string
+	Folder        string
+	Shell         string
+	ShellVersion  string
+	UserName      string
+	HostName      string
+	Code          int
+	Env           map[string]string
+	Var           SimpleMap
+	OS            string
+	WSL           bool
+	PromptCount   int
+	SHLVL         int
+	Segments      *ConcurrentMap
+	SegmentsCache SimpleMap
 
 	initialized bool
 	sync.RWMutex
@@ -745,7 +746,9 @@ func (env *Shell) saveTemplateCache() {
 	if !canSave {
 		return
 	}
-	templateCache, err := json.Marshal(env.TemplateCache())
+	cache := env.TemplateCache()
+	cache.SegmentsCache = cache.Segments.SimpleMap()
+	templateCache, err := json.Marshal(cache)
 	if err == nil {
 		env.fileCache.Set(TEMPLATECACHE, string(templateCache), 1440)
 	}
@@ -769,6 +772,7 @@ func (env *Shell) LoadTemplateCache() {
 		env.Error(err)
 		return
 	}
+	tmplCache.Segments = tmplCache.SegmentsCache.ConcurrentMap()
 	tmplCache.initialized = true
 	env.tmplCache = &tmplCache
 }
diff --git a/src/template/text.go b/src/template/text.go
index 20f4ea28..9bee237d 100644
--- a/src/template/text.go
+++ b/src/template/text.go
@@ -163,7 +163,7 @@ func (t *Text) cleanTemplate() {
 				// as we can't provide a clean way to access the list
 				// of segments, we need to replace the property with
 				// the list of segments so they can be accessed directly
-				property = strings.Replace(property, ".Segments", ".Segments.List", 1)
+				property = strings.Replace(property, ".Segments", ".Segments.SimpleMap", 1)
 				result += property
 			} else {
 				// check if we have the same property in Data
diff --git a/src/template/text_test.go b/src/template/text_test.go
index 27f3ee83..b0d4d428 100644
--- a/src/template/text_test.go
+++ b/src/template/text_test.go
@@ -327,7 +327,7 @@ func TestCleanTemplate(t *testing.T) {
 		},
 		{
 			Case:     "Replace a direct call to .Segments with .Segments.List",
-			Expected: `{{.Segments.List.Git.Repo}}`,
+			Expected: `{{.Segments.SimpleMap.Git.Repo}}`,
 			Template: `{{.Segments.Git.Repo}}`,
 		},
 	}