From bf7d1e89ddd1c2aac7534edd30fc700e45b7f73a Mon Sep 17 00:00:00 2001
From: Tristan Swadell <tswadell@google.com>
Date: Tue, 4 Apr 2023 12:03:33 -0700
Subject: [PATCH] Support for optional on unparse (#673)

---
 parser/unparser.go      | 53 +++++++++++++++++++++++++++++++++++------
 parser/unparser_test.go |  6 +++++
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/parser/unparser.go b/parser/unparser.go
index ea20efc9..5ff979aa 100644
--- a/parser/unparser.go
+++ b/parser/unparser.go
@@ -106,9 +106,15 @@ func (un *unparser) visitCall(expr *exprpb.Expr) error {
 	// ternary operator
 	case operators.Conditional:
 		return un.visitCallConditional(expr)
+	// optional select operator
+	case operators.OptSelect:
+		return un.visitOptSelect(expr)
 	// index operator
 	case operators.Index:
 		return un.visitCallIndex(expr)
+	// optional index operator
+	case operators.OptIndex:
+		return un.visitCallOptIndex(expr)
 	// unary operators
 	case operators.LogicalNot, operators.Negate:
 		return un.visitCallUnary(expr)
@@ -218,6 +224,14 @@ func (un *unparser) visitCallFunc(expr *exprpb.Expr) error {
 }
 
 func (un *unparser) visitCallIndex(expr *exprpb.Expr) error {
+	return un.visitCallIndexInternal(expr, "[")
+}
+
+func (un *unparser) visitCallOptIndex(expr *exprpb.Expr) error {
+	return un.visitCallIndexInternal(expr, "[?")
+}
+
+func (un *unparser) visitCallIndexInternal(expr *exprpb.Expr, op string) error {
 	c := expr.GetCallExpr()
 	args := c.GetArgs()
 	nested := isBinaryOrTernaryOperator(args[0])
@@ -225,7 +239,7 @@ func (un *unparser) visitCallIndex(expr *exprpb.Expr) error {
 	if err != nil {
 		return err
 	}
-	un.str.WriteString("[")
+	un.str.WriteString(op)
 	err = un.visit(args[1])
 	if err != nil {
 		return err
@@ -289,8 +303,15 @@ func (un *unparser) visitIdent(expr *exprpb.Expr) error {
 func (un *unparser) visitList(expr *exprpb.Expr) error {
 	l := expr.GetListExpr()
 	elems := l.GetElements()
+	optIndices := make(map[int]bool, len(elems))
+	for _, idx := range l.GetOptionalIndices() {
+		optIndices[int(idx)] = true
+	}
 	un.str.WriteString("[")
 	for i, elem := range elems {
+		if optIndices[i] {
+			un.str.WriteString("?")
+		}
 		err := un.visit(elem)
 		if err != nil {
 			return err
@@ -303,20 +324,32 @@ func (un *unparser) visitList(expr *exprpb.Expr) error {
 	return nil
 }
 
+func (un *unparser) visitOptSelect(expr *exprpb.Expr) error {
+	c := expr.GetCallExpr()
+	args := c.GetArgs()
+	operand := args[0]
+	field := args[1].GetConstExpr().GetStringValue()
+	return un.visitSelectInternal(operand, false, ".?", field)
+}
+
 func (un *unparser) visitSelect(expr *exprpb.Expr) error {
 	sel := expr.GetSelectExpr()
+	return un.visitSelectInternal(sel.GetOperand(), sel.GetTestOnly(), ".", sel.GetField())
+}
+
+func (un *unparser) visitSelectInternal(operand *exprpb.Expr, testOnly bool, op string, field string) error {
 	// handle the case when the select expression was generated by the has() macro.
-	if sel.GetTestOnly() {
+	if testOnly {
 		un.str.WriteString("has(")
 	}
-	nested := !sel.GetTestOnly() && isBinaryOrTernaryOperator(sel.GetOperand())
-	err := un.visitMaybeNested(sel.GetOperand(), nested)
+	nested := !testOnly && isBinaryOrTernaryOperator(operand)
+	err := un.visitMaybeNested(operand, nested)
 	if err != nil {
 		return err
 	}
-	un.str.WriteString(".")
-	un.str.WriteString(sel.GetField())
-	if sel.GetTestOnly() {
+	un.str.WriteString(op)
+	un.str.WriteString(field)
+	if testOnly {
 		un.str.WriteString(")")
 	}
 	return nil
@@ -339,6 +372,9 @@ func (un *unparser) visitStructMsg(expr *exprpb.Expr) error {
 	un.str.WriteString("{")
 	for i, entry := range entries {
 		f := entry.GetFieldKey()
+		if entry.GetOptionalEntry() {
+			un.str.WriteString("?")
+		}
 		un.str.WriteString(f)
 		un.str.WriteString(": ")
 		v := entry.GetValue()
@@ -360,6 +396,9 @@ func (un *unparser) visitStructMap(expr *exprpb.Expr) error {
 	un.str.WriteString("{")
 	for i, entry := range entries {
 		k := entry.GetMapKey()
+		if entry.GetOptionalEntry() {
+			un.str.WriteString("?")
+		}
 		err := un.visit(k)
 		if err != nil {
 			return err
diff --git a/parser/unparser_test.go b/parser/unparser_test.go
index fad688e9..7696d724 100644
--- a/parser/unparser_test.go
+++ b/parser/unparser_test.go
@@ -99,6 +99,11 @@ func TestUnparse(t *testing.T) {
 		{name: "cond_binop", in: `(x < 5) ? x : 5`},
 		{name: "cond_binop_binop", in: `(x > 5) ? (x - 5) : 0`},
 		{name: "cond_cond_binop", in: `(x > 5) ? ((x > 10) ? (x - 10) : 5) : 0`},
+		{name: "select_opt", in: `a.?b`},
+		{name: "index_opt", in: `a[?b]`},
+		{name: "list_lit_opt", in: `[?a, ?b, c]`},
+		{name: "map_lit_opt", in: `{?a: b, c: d}`},
+		{name: "msg_fields_opt", in: `v1alpha1.Expr{?id: id, call_expr: v1alpha1.Call_Expr{function: "name"}}`},
 
 		// Equivalent expressions form unparse which do not match the originals.
 		{name: "call_add_equiv", in: `a+b-c`, out: `a + b - c`},
@@ -424,6 +429,7 @@ func TestUnparse(t *testing.T) {
 			prsr, err := NewParser(
 				Macros(AllMacros...),
 				PopulateMacroCalls(tc.requiresMacroCalls),
+				EnableOptionalSyntax(true),
 			)
 			if err != nil {
 				t.Fatalf("NewParser() failed: %v", err)