From d2daabfa022c7abd5575f5a8a4c21e600b5ba58d Mon Sep 17 00:00:00 2001
From: Cole Faust <colefaust@google.com>
Date: Mon, 12 Dec 2022 17:38:01 -0800
Subject: [PATCH] Add support for and/or in mk2rbc

Bug: 262303006
Test: go test
Change-Id: I04130648084775c0828e95cd18a82e348c3f09eb
---
 mk2rbc/mk2rbc.go      | 47 +++++++++++++++++++++++++++++++++++++++++++
 mk2rbc/mk2rbc_test.go | 32 +++++++++++++++++++++++++++++
 mk2rbc/types.go       | 21 +++++++++++++++++++
 3 files changed, 100 insertions(+)

diff --git a/mk2rbc/mk2rbc.go b/mk2rbc/mk2rbc.go
index aa48e6377..8b8629b94 100644
--- a/mk2rbc/mk2rbc.go
+++ b/mk2rbc/mk2rbc.go
@@ -77,6 +77,7 @@ var knownFunctions = map[string]interface {
 	"add-to-product-copy-files-if-exists":  &simpleCallParser{name: baseName + ".copy_if_exists", returnType: starlarkTypeList},
 	"addprefix":                            &simpleCallParser{name: baseName + ".addprefix", returnType: starlarkTypeList},
 	"addsuffix":                            &simpleCallParser{name: baseName + ".addsuffix", returnType: starlarkTypeList},
+	"and":                                  &andOrParser{isAnd: true},
 	"copy-files":                           &simpleCallParser{name: baseName + ".copy_files", returnType: starlarkTypeList},
 	"dir":                                  &simpleCallParser{name: baseName + ".dir", returnType: starlarkTypeString},
 	"dist-for-goals":                       &simpleCallParser{name: baseName + ".mkdist_for_goals", returnType: starlarkTypeVoid, addGlobals: true},
@@ -105,6 +106,7 @@ var knownFunctions = map[string]interface {
 	"math_gt":                              &mathComparisonCallParser{op: ">"},
 	"math_lt":                              &mathComparisonCallParser{op: "<"},
 	"my-dir":                               &myDirCallParser{},
+	"or":                                   &andOrParser{isAnd: false},
 	"patsubst":                             &substCallParser{fname: "patsubst"},
 	"product-copy-files-by-pattern":        &simpleCallParser{name: baseName + ".product_copy_files_by_pattern", returnType: starlarkTypeList},
 	"require-artifacts-in-path":            &simpleCallParser{name: baseName + ".require_artifacts_in_path", returnType: starlarkTypeVoid, addHandle: true},
@@ -1430,6 +1432,51 @@ func (p *myDirCallParser) parse(ctx *parseContext, node mkparser.Node, args *mkp
 	return &stringLiteralExpr{literal: filepath.Dir(ctx.script.mkFile)}
 }
 
+type andOrParser struct {
+	isAnd bool
+}
+
+func (p *andOrParser) parse(ctx *parseContext, node mkparser.Node, args *mkparser.MakeString) starlarkExpr {
+	if args.Empty() {
+		return ctx.newBadExpr(node, "and/or function must have at least 1 argument")
+	}
+	op := "or"
+	if p.isAnd {
+		op = "and"
+	}
+
+	argsParsed := make([]starlarkExpr, 0)
+
+	for _, arg := range args.Split(",") {
+		arg.TrimLeftSpaces()
+		arg.TrimRightSpaces()
+		x := ctx.parseMakeString(node, arg)
+		if xBad, ok := x.(*badExpr); ok {
+			return xBad
+		}
+		argsParsed = append(argsParsed, x)
+	}
+	typ := starlarkTypeUnknown
+	for _, arg := range argsParsed {
+		if typ != arg.typ() && arg.typ() != starlarkTypeUnknown && typ != starlarkTypeUnknown {
+			return ctx.newBadExpr(node, "Expected all arguments to $(or) or $(and) to have the same type, found %q and %q", typ.String(), arg.typ().String())
+		}
+		if arg.typ() != starlarkTypeUnknown {
+			typ = arg.typ()
+		}
+	}
+	result := argsParsed[0]
+	for _, arg := range argsParsed[1:] {
+		result = &binaryOpExpr{
+			left:       result,
+			right:      arg,
+			op:         op,
+			returnType: typ,
+		}
+	}
+	return result
+}
+
 type isProductInListCallParser struct{}
 
 func (p *isProductInListCallParser) parse(ctx *parseContext, node mkparser.Node, args *mkparser.MakeString) starlarkExpr {
diff --git a/mk2rbc/mk2rbc_test.go b/mk2rbc/mk2rbc_test.go
index 31555d3f5..7e14e259d 100644
--- a/mk2rbc/mk2rbc_test.go
+++ b/mk2rbc/mk2rbc_test.go
@@ -1627,6 +1627,38 @@ def init(g, handle):
   g["MY_VAR_3"] = (cfg).get(g["MY_VAR_2"], (g).get(g["MY_VAR_2"], ""))
   g["MY_VAR_4"] = rblf.mk2rbc_error("product.mk:5", "cannot handle invoking foo")
   g["MY_VAR_5"] = rblf.mk2rbc_error("product.mk:6", "reference is too complex: $(MY_VAR_2) bar")
+`,
+	},
+	{
+		desc:   "Conditional functions",
+		mkname: "product.mk",
+		in: `
+B := foo
+X := $(or $(A))
+X := $(or $(A),$(B))
+X := $(or $(A),$(B),$(C))
+X := $(and $(A))
+X := $(and $(A),$(B))
+X := $(and $(A),$(B),$(C))
+X := $(or $(A),$(B)) Y
+
+D := $(wildcard *.mk)
+X := $(or $(B),$(D))
+`,
+		expected: `load("//build/make/core:product_config.rbc", "rblf")
+
+def init(g, handle):
+  cfg = rblf.cfg(handle)
+  g["B"] = "foo"
+  g["X"] = g.get("A", "")
+  g["X"] = g.get("A", "") or g["B"]
+  g["X"] = g.get("A", "") or g["B"] or g.get("C", "")
+  g["X"] = g.get("A", "")
+  g["X"] = g.get("A", "") and g["B"]
+  g["X"] = g.get("A", "") and g["B"] and g.get("C", "")
+  g["X"] = "%s Y" % g.get("A", "") or g["B"]
+  g["D"] = rblf.expand_wildcard("*.mk")
+  g["X"] = rblf.mk2rbc_error("product.mk:12", "Expected all arguments to $(or) or $(and) to have the same type, found \"string\" and \"list\"")
 `,
 	},
 }
diff --git a/mk2rbc/types.go b/mk2rbc/types.go
index 46c6aa9b3..ac3250722 100644
--- a/mk2rbc/types.go
+++ b/mk2rbc/types.go
@@ -14,6 +14,8 @@
 
 package mk2rbc
 
+import "fmt"
+
 // Starlark expression types we use
 type starlarkType int
 
@@ -31,6 +33,25 @@ const (
 	starlarkTypeVoid    starlarkType = iota
 )
 
+func (t starlarkType) String() string {
+	switch t {
+	case starlarkTypeList:
+		return "list"
+	case starlarkTypeString:
+		return "string"
+	case starlarkTypeInt:
+		return "int"
+	case starlarkTypeBool:
+		return "bool"
+	case starlarkTypeVoid:
+		return "void"
+	case starlarkTypeUnknown:
+		return "unknown"
+	default:
+		panic(fmt.Sprintf("Unknown starlark type %d", t))
+	}
+}
+
 type hiddenArgType int
 
 const (