From daed7465e55570c5ca11e10ef03bcb8a248b1743 Mon Sep 17 00:00:00 2001
From: baileyspencer <bailey.spencer@rainforestpay.com>
Date: Thu, 24 Oct 2024 09:25:38 -0400
Subject: [PATCH] Add OrderedDict. Refactor Dict.render to utilize OrderedDict

---
 jen/dict.go          | 45 +++++++++--------------------
 jen/examples_test.go | 20 +++++++++++++
 jen/ordered_dict.go  | 69 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 32 deletions(-)
 create mode 100644 jen/ordered_dict.go

diff --git a/jen/dict.go b/jen/dict.go
index 260cf92..a3d98a1 100644
--- a/jen/dict.go
+++ b/jen/dict.go
@@ -19,51 +19,32 @@ func DictFunc(f func(Dict)) Dict {
 }
 
 func (d Dict) render(f *File, w io.Writer, s *Statement) error {
-	first := true
-	// must order keys to ensure repeatable source
-	type kv struct {
-		k Code
-		v Code
-	}
-	lookup := map[string]kv{}
-	keys := []string{}
+	lookup := make(map[string]kv, len(d))
+	keys := make([]string, 0, len(d))
+
+	buf := &bytes.Buffer{}
 	for k, v := range d {
 		if k.isNull(f) || v.isNull(f) {
 			continue
 		}
-		buf := &bytes.Buffer{}
 		if err := k.render(f, buf, nil); err != nil {
 			return err
 		}
 		keys = append(keys, buf.String())
 		lookup[buf.String()] = kv{k: k, v: v}
+		buf.Reset()
 	}
+
+	// must order keys to ensure repeatable source
 	sort.Strings(keys)
+
+	ordered := make([]kv, 0, len(keys))
 	for _, key := range keys {
-		k := lookup[key].k
-		v := lookup[key].v
-		if first && len(keys) > 1 {
-			if _, err := w.Write([]byte("\n")); err != nil {
-				return err
-			}
-			first = false
-		}
-		if err := k.render(f, w, nil); err != nil {
-			return err
-		}
-		if _, err := w.Write([]byte(":")); err != nil {
-			return err
-		}
-		if err := v.render(f, w, nil); err != nil {
-			return err
-		}
-		if len(keys) > 1 {
-			if _, err := w.Write([]byte(",\n")); err != nil {
-				return err
-			}
-		}
+		ordered = append(ordered, kv{k: lookup[key].k, v: lookup[key].v})
 	}
-	return nil
+
+	dict := &OrderedDict{items: ordered}
+	return dict.render(f, w, s)
 }
 
 func (d Dict) isNull(f *File) bool {
diff --git a/jen/examples_test.go b/jen/examples_test.go
index f1d825a..4990cf9 100644
--- a/jen/examples_test.go
+++ b/jen/examples_test.go
@@ -1577,6 +1577,26 @@ func ExampleDictFunc() {
 	// }
 }
 
+func ExampleOrderedDictFunc() {
+	c := Id("a").Op(":=").Map(String()).String().Values(OrderedDictFunc(func(o *OrderedDict) {
+		o.Add(Lit("c"), Lit("d"))
+		o.Add(Lit("a"), Lit("b"))
+	}))
+	fmt.Printf("%#v", c)
+	// Output:
+	// a := map[string]string{
+	// 	"c": "d",
+	// 	"a": "b",
+	// }
+}
+
+func ExampleOrderedDictFunc_nil() {
+	c := Id("a").Op(":=").Map(String()).String().Values(OrderedDictFunc(func(o *OrderedDict) {}))
+	fmt.Printf("%#v", c)
+	// Output:
+	// a := map[string]string{}
+}
+
 func ExampleDefs() {
 	c := Const().Defs(
 		Id("a").Op("=").Lit("a"),
diff --git a/jen/ordered_dict.go b/jen/ordered_dict.go
new file mode 100644
index 0000000..dd74322
--- /dev/null
+++ b/jen/ordered_dict.go
@@ -0,0 +1,69 @@
+package jen
+
+import "io"
+
+// OrderedDictFunc executes a func(*OrderedDict) to generate the value. Use OrderedDict.Add to append key/value pairs.
+// Use with Values for map or composite literals.
+func OrderedDictFunc(f func(*OrderedDict)) *OrderedDict {
+	o := &OrderedDict{}
+	f(o)
+	return o
+}
+
+// OrderedDict renders as key/value pairs. Use with Values for map or composite
+// literals.
+type OrderedDict struct {
+	items []kv
+}
+
+type kv struct {
+	k Code
+	v Code
+}
+
+// Add appends a key/value pairs.
+func (o *OrderedDict) Add(key Code, value Code) *OrderedDict {
+	o.items = append(o.items, kv{k: key, v: value})
+	return o
+}
+
+func (o *OrderedDict) render(f *File, w io.Writer, _ *Statement) error {
+	first := true
+	for _, item := range o.items {
+		if first && len(o.items) > 1 {
+			if _, err := w.Write([]byte("\n")); err != nil {
+				return err
+			}
+			first = false
+		}
+		if err := item.k.render(f, w, nil); err != nil {
+			return err
+		}
+		if _, err := w.Write([]byte(":")); err != nil {
+			return err
+		}
+		if err := item.v.render(f, w, nil); err != nil {
+			return err
+		}
+		if len(o.items) > 1 {
+			if _, err := w.Write([]byte(",\n")); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (o *OrderedDict) isNull(f *File) bool {
+	if len(o.items) == 0 {
+		return true
+	}
+	for _, item := range o.items {
+		if !item.k.isNull(f) && !item.v.isNull(f) {
+			// if any of the key/value pairs are both not null, the OrderedDict is not
+			// null
+			return false
+		}
+	}
+	return true
+}