From ea610029b2f3beef372b3764a55658ca95fa4e23 Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Thu, 3 Aug 2023 19:36:35 +0200
Subject: [PATCH] fix(engine): use concurrent map for Segments

resolves #4116
---
 src/platform/concurrent_map.go |  9 +++++++++
 src/platform/shell.go          | 19 ++++---------------
 src/template/text.go           |  6 ++++++
 src/template/text_test.go      | 18 ++++++++++++++----
 4 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/src/platform/concurrent_map.go b/src/platform/concurrent_map.go
index 696c8345..45c446eb 100644
--- a/src/platform/concurrent_map.go
+++ b/src/platform/concurrent_map.go
@@ -37,3 +37,12 @@ func (c *ConcurrentMap) Delete(key string) {
 func (c *ConcurrentMap) List() map[string]interface{} {
 	return c.values
 }
+
+func (c *ConcurrentMap) Contains(key string) bool {
+	c.RLock()
+	defer c.RUnlock()
+	if _, ok := c.values[key]; ok {
+		return true
+	}
+	return false
+}
diff --git a/src/platform/shell.go b/src/platform/shell.go
index 627e9c93..9bcf88c1 100644
--- a/src/platform/shell.go
+++ b/src/platform/shell.go
@@ -166,13 +166,6 @@ type SystemInfo struct {
 	Disks map[string]disk.IOCountersStat
 }
 
-type SegmentsCache map[string]interface{}
-
-func (s *SegmentsCache) Contains(key string) bool {
-	_, ok := (*s)[key]
-	return ok
-}
-
 type TemplateCache struct {
 	Root         bool
 	PWD          string
@@ -188,22 +181,18 @@ type TemplateCache struct {
 	WSL          bool
 	PromptCount  int
 	SHLVL        int
-	Segments     SegmentsCache
+	Segments     *ConcurrentMap
 
 	initialized bool
 	sync.RWMutex
 }
 
 func (t *TemplateCache) AddSegmentData(key string, value interface{}) {
-	t.Lock()
-	t.Segments[key] = value
-	t.Unlock()
+	t.Segments.Set(key, value)
 }
 
 func (t *TemplateCache) RemoveSegmentData(key string) {
-	t.Lock()
-	delete(t.Segments, key)
-	t.Unlock()
+	t.Segments.Delete(key)
 }
 
 type Environment interface {
@@ -802,7 +791,7 @@ func (env *Shell) TemplateCache() *TemplateCache {
 	tmplCache.ShellVersion = env.CmdFlags.ShellVersion
 	tmplCache.Code, _ = env.StatusCodes()
 	tmplCache.WSL = env.IsWsl()
-	tmplCache.Segments = make(map[string]interface{})
+	tmplCache.Segments = NewConcurrentMap()
 	tmplCache.PromptCount = env.CmdFlags.PromptCount
 	tmplCache.Env = make(map[string]string)
 	tmplCache.Var = make(map[string]interface{})
diff --git a/src/template/text.go b/src/template/text.go
index a4bf9b64..20f4ea28 100644
--- a/src/template/text.go
+++ b/src/template/text.go
@@ -159,6 +159,12 @@ func (t *Text) cleanTemplate() {
 			// end of a variable, needs to be appended
 			if !isKnownVariable(property) {
 				result += ".Data" + property
+			} else if strings.HasPrefix(property, ".Segments") && !strings.HasSuffix(property, ".Contains") {
+				// 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)
+				result += property
 			} else {
 				// check if we have the same property in Data
 				// and replace it with the Data property so it
diff --git a/src/template/text_test.go b/src/template/text_test.go
index 29074d3d..27f3ee83 100644
--- a/src/template/text_test.go
+++ b/src/template/text_test.go
@@ -320,6 +320,16 @@ func TestCleanTemplate(t *testing.T) {
 			Expected: "{{.Data.OS}}",
 			Template: "{{.OS}}",
 		},
+		{
+			Case:     "Keep .Contains intact for Segments",
+			Expected: `{{.Segments.Contains "Git"}}`,
+			Template: `{{.Segments.Contains "Git"}}`,
+		},
+		{
+			Case:     "Replace a direct call to .Segments with .Segments.List",
+			Expected: `{{.Segments.List.Git.Repo}}`,
+			Template: `{{.Segments.Git.Repo}}`,
+		},
 	}
 	for _, tc := range cases {
 		tmpl := &Text{
@@ -342,11 +352,11 @@ func TestSegmentContains(t *testing.T) {
 	}
 
 	env := &mock.MockedEnvironment{}
+	segments := platform.NewConcurrentMap()
+	segments.Set("Git", "foo")
 	env.On("TemplateCache").Return(&platform.TemplateCache{
-		Env: make(map[string]string),
-		Segments: map[string]interface{}{
-			"Git": nil,
-		},
+		Env:      make(map[string]string),
+		Segments: segments,
 	})
 	for _, tc := range cases {
 		tmpl := &Text{