From f86a886abd5081305596f9ff4993a21d92bde220 Mon Sep 17 00:00:00 2001
From: Tristan Swadell <tswadell@google.com>
Date: Thu, 25 Mar 2021 12:11:52 -0700
Subject: [PATCH] Fix infinite loop in AST prune (#419)

---
 interpreter/prune.go      |  9 ++++-----
 interpreter/prune_test.go | 24 ++++++++++++++++++++++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/interpreter/prune.go b/interpreter/prune.go
index 2ad522bc..7bcbb6c9 100644
--- a/interpreter/prune.go
+++ b/interpreter/prune.go
@@ -382,12 +382,11 @@ func (p *astPruner) existsWithKnownValue(id int64) bool {
 func (p *astPruner) nextID() int64 {
 	for {
 		_, found := p.state.Value(p.nextExprID)
-		if found {
-			break
+		if !found {
+			next := p.nextExprID
+			p.nextExprID++
+			return next
 		}
 		p.nextExprID++
 	}
-	next := p.nextExprID
-	p.nextExprID++
-	return next
 }
diff --git a/interpreter/prune_test.go b/interpreter/prune_test.go
index 32929e28..e2940ee7 100644
--- a/interpreter/prune_test.go
+++ b/interpreter/prune_test.go
@@ -24,12 +24,19 @@ import (
 )
 
 type testInfo struct {
-	in   Activation
+	in   interface{}
 	expr string
 	out  string
 }
 
 var testCases = []testInfo{
+	{
+		in: map[string]interface{}{
+			"msg": map[string]string{"foo": "bar"},
+		},
+		expr: `msg`,
+		out:  `{"foo": "bar"}`,
+	},
 	{
 		expr: `true && false`,
 		out:  `false`,
@@ -128,10 +135,11 @@ func TestPrune(t *testing.T) {
 		reg := newTestRegistry(t)
 		attrs := NewPartialAttributeFactory(containers.DefaultContainer, reg, reg)
 		interp := NewStandardInterpreter(containers.DefaultContainer, reg, reg, attrs)
+
 		interpretable, _ := interp.NewUncheckedInterpretable(
 			ast.Expr,
 			ExhaustiveEval(state))
-		interpretable.Eval(tst.in)
+		interpretable.Eval(testActivation(t, tst.in))
 		newExpr := PruneAst(ast.Expr, state)
 		actual, err := parser.Unparse(newExpr, nil)
 		if err != nil {
@@ -151,3 +159,15 @@ func unknownActivation(vars ...string) PartialActivation {
 	a, _ := NewPartialActivation(map[string]interface{}{}, pats...)
 	return a
 }
+
+func testActivation(t *testing.T, in interface{}) Activation {
+	t.Helper()
+	if in == nil {
+		return EmptyActivation()
+	}
+	a, err := NewActivation(in)
+	if err != nil {
+		t.Fatalf("NewActivation(%v) failed: %v", in, err)
+	}
+	return a
+}