Add (*TemplateString).Copy() method (#8159)

Also:
- Add `(*TemplateString).Equal()` because why not.
- Update `x.Compare(y) == 0` to instead use `x.Equal(y)` where possible

Fixes #8158

Signed-off-by: Anders Eknert <anders.eknert@apple.com>

internal/planner/planner.go

@@ -1768,7 +1768,7 @@ func (p *Planner) planRef(ref ast.Ref, iter planiter) error {
 		return errors.New("illegal ref: non-var head")
 	}
 
-	if head.Compare(ast.DefaultRootDocument.Value) == 0 {
+	if head.Equal(ast.DefaultRootDocument.Value) {
 		virtual := p.rules.Get(ref[0].Value)
 		base := &baseptr{local: p.vars.GetOrEmpty(ast.DefaultRootDocument.Value.(ast.Var))}
 		return p.planRefData(virtual, base, ref, 1, iter)

v1/ast/compare.go

@@ -327,14 +327,6 @@ func TermValueEqual(a, b *Term) bool {
 }
 
 func ValueEqual(a, b Value) bool {
-	// TODO(ae): why doesn't this work the same?
-	//
-	// case interface{ Equal(Value) bool }:
-	// 	   return v.Equal(b)
-	//
-	// When put on top, golangci-lint even flags the other cases as unreachable..
-	// but TestTopdownVirtualCache will have failing test cases when we replace
-	// the other cases with the above one.. 🤔
 	switch v := a.(type) {
 	case Null:
 		return v.Equal(b)
@@ -350,6 +342,8 @@ func ValueEqual(a, b Value) bool {
 		return v.Equal(b)
 	case *Array:
 		return v.Equal(b)
+	case *TemplateString:
+		return v.Equal(b)
 	}
 
 	return a.Compare(b) == 0

v1/ast/compile.go

@@ -3746,7 +3746,7 @@ func NewModuleTree(mods map[string]*Module) *ModuleTreeNode {
 			c, ok := node.Children[x.Value]
 			if !ok {
 				var hide bool
-				if i == 1 && x.Value.Compare(SystemDocumentKey) == 0 {
+				if i == 1 && SystemDocumentKey.Equal(x.Value) {
 					hide = true
 				}
 				c = &ModuleTreeNode{

v1/ast/compile_test.go

@@ -12943,3 +12943,27 @@ func TestCompilerInitWithDefaultModuleLoader(t *testing.T) {
 		t.Error("expected bar.rego from defaultModuleLoader in result")
 	}
 }
+
+// Verify fix for https://github.com/open-policy-agent/opa/issues/8158
+func TestCompilerCopiesTemplateStrings(t *testing.T) {
+	mod := MustParseModule(`package p
+	s contains z if {
+		some y in [1, 2, 3]
+		z := $"{y} "
+	}`)
+	cpy := mod.Copy()
+
+	c1 := NewCompiler()
+	if c1.Compile(map[string]*Module{"p.rego": mod}); c1.Failed() {
+		t.Fatalf("unexpected compile errors: %v", c1.Errors)
+	}
+
+	c2 := NewCompiler()
+	if c2.Compile(map[string]*Module{"p.rego": mod}); c2.Failed() {
+		t.Fatalf("unexpected compile errors: %v", c2.Errors)
+	}
+
+	if !mod.Equal(cpy) {
+		t.Fatalf("expected module to be unchanged after compilation")
+	}
+}

v1/ast/index.go

@@ -412,7 +412,7 @@ func (i *refindices) updateGlobMatch(rule *Rule, expr *Expr) {
 		if _, ok := match.Value.(Var); ok {
 			var ref Ref
 			for _, other := range i.rules[rule] {
-				if _, ok := other.Value.(Var); ok && other.Value.Compare(match.Value) == 0 {
+				if ov, ok := other.Value.(Var); ok && ov.Equal(match.Value) {
 					ref = other.Ref
 				}
 			}

v1/ast/oracle/matcher.go

@@ -190,7 +190,7 @@ func termMatchesVar(t *ast.Term, name ast.Var) bool {
 
 	v, ok := t.Value.(ast.Var)
 
-	return ok && v.Compare(name) == 0
+	return ok && v.Equal(name)
 }
 
 // findRulesDefinition looks up rules for a given ref. Rules appear in various

v1/ast/term.go

@@ -43,6 +43,7 @@ func NewLocation(text []byte, file string, row int, col int) *Location {
 // - Variables, References
 // - Array, Set, and Object Comprehensions
 // - Calls
+// - Template Strings
 type Value interface {
 	Compare(other Value) int      // Compare returns <0, 0, or >0 if this Value is less than, equal to, or greater than other, respectively.
 	Find(path Ref) (Value, error) // Find returns value referred to by path or an error if path is not found.
@@ -351,6 +352,8 @@ func (term *Term) Copy() *Term {
 		cpy.Value = v.Copy()
 	case *SetComprehension:
 		cpy.Value = v.Copy()
+	case *TemplateString:
+		cpy.Value = v.Copy()
 	case Call:
 		cpy.Value = v.Copy()
 	}
@@ -833,6 +836,40 @@ type TemplateString struct {
 	MultiLine bool   `json:"multi_line"`
 }
 
+func (ts *TemplateString) Copy() *TemplateString {
+	cpy := &TemplateString{MultiLine: ts.MultiLine, Parts: make([]Node, len(ts.Parts))}
+	for i, p := range ts.Parts {
+		switch v := p.(type) {
+		case *Expr:
+			cpy.Parts[i] = v.Copy()
+		case *Term:
+			cpy.Parts[i] = v.Copy()
+		}
+	}
+	return cpy
+}
+
+func (ts *TemplateString) Equal(other Value) bool {
+	if o, ok := other.(*TemplateString); ok && ts.MultiLine == o.MultiLine && len(ts.Parts) == len(o.Parts) {
+		for i, p := range ts.Parts {
+			switch v := p.(type) {
+			case *Expr:
+				if ope, ok := o.Parts[i].(*Expr); !ok || !v.Equal(ope) {
+					return false
+				}
+			case *Term:
+				if opt, ok := o.Parts[i].(*Term); !ok || !v.Equal(opt) {
+					return false
+				}
+			default:
+				return false
+			}
+		}
+		return true
+	}
+	return false
+}
+
 func (ts *TemplateString) Compare(other Value) int {
 	if ots, ok := other.(*TemplateString); ok {
 		if ts.MultiLine != ots.MultiLine {

v1/ast/term_test.go

@@ -1689,6 +1689,20 @@ func TestLazyObjectCompare(t *testing.T) {
 	assertForced(t, x, true)
 }
 
+func TestTemplateStringEqual(t *testing.T) {
+	a := MustParseTerm(`$"hello {world}!"`).Value.(*TemplateString)
+	b := MustParseTerm(`$"hello {world}!"`).Value.(*TemplateString)
+	c := MustParseTerm(`$"goodbye {world}!"`).Value.(*TemplateString)
+
+	if !a.Equal(b) {
+		t.Errorf("Expected %v to equal %v", a, b)
+	}
+
+	if a.Equal(c) {
+		t.Errorf("Expected %v to not equal %v", a, c)
+	}
+}
+
 func assertForced(t *testing.T, x Object, forced bool) {
 	t.Helper()
 	l, ok := x.(*lazyObj)

v1/ast/transform_test.go

@@ -121,7 +121,7 @@ func TestTransformRefsAndRuleHeads(t *testing.T) {
 p.q.this.fo[x] = y if { x := "x"; y := "y" }`)
 
 	result, err := TransformRefs(module, func(r Ref) (Value, error) {
-		if r[0].Value.Compare(Var("p")) == 0 {
+		if Var("p").Equal(r[0].Value) {
 			r[2] = StringTerm("that")
 		}
 		return r, nil

v1/repl/repl.go

@@ -930,7 +930,7 @@ func (r *REPL) parserOptions() (ast.ParserOptions, error) {
 		opts, err := future.ParserOptionsFromFutureImports(r.modules[r.currentModuleID].Imports)
 		if err == nil {
 			for _, i := range r.modules[r.currentModuleID].Imports {
-				if ast.Compare(i.Path.Value, ast.RegoV1CompatibleRef) == 0 {
+				if ast.RegoV1CompatibleRef.Equal(i.Path.Value) {
 					opts.RegoVersion = ast.RegoV1
 				}
 			}

v1/tester/runner.go

@@ -799,7 +799,7 @@ func injectTestCaseFunc(compiler *ast.Compiler) *ast.Error {
 					expr := rule.Body[i]
 
 					ast.WalkVars(expr, func(v ast.Var) bool {
-						if term.Value.Compare(v) == 0 {
+						if v.Equal(term.Value) {
 							injectBelowMap.Put(v, ast.Number(strconv.Itoa(i)))
 						}
 						return false

v1/topdown/copypropagation/copypropagation.go

@@ -344,7 +344,7 @@ func (p *CopyPropagator) livevarRef(a *ast.Term) bool {
 	}
 
 	for _, v := range p.sorted {
-		if ref[0].Value.Compare(v) == 0 {
+		if v.Equal(ref[0].Value) {
 			return true
 		}
 	}
@@ -403,7 +403,7 @@ func containedIn(value ast.Value, x any) bool {
 			if v, ok := value.(ast.Ref); ok {
 				match = x.HasPrefix(v)
 			} else {
-				match = x.Compare(value) == 0
+				match = x.Equal(value)
 			}
 			if stop || match {
 				stop = true

v1/topdown/tokens.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"hash"
 	"math/big"
+	"strconv"
 	"strings"
 
 	"github.com/lestrrat-go/jwx/v3/jwk"
@@ -1131,8 +1132,8 @@ func builtinJWTDecodeVerify(bctx BuiltinContext, operands []*ast.Term, iter func
 		switch v := nbf.Value.(type) {
 		case ast.Number:
 			// constraints.time is in nanoseconds but nbf Value is in seconds
-			compareTime := ast.FloatNumberTerm(constraints.time / 1000000000)
-			if ast.Compare(compareTime, v) == -1 {
+			compareTime := ast.Number(strconv.FormatFloat(constraints.time/1000000000, 'g', -1, 64))
+			if compareTime.Compare(v) == -1 {
 				return iter(unverified)
 			}
 		default: