From 9df1d736572536445794e327af5925b0bc097f77 Mon Sep 17 00:00:00 2001
From: Cole Faust <colefaust@google.com>
Date: Tue, 26 Apr 2022 16:27:22 -0700
Subject: [PATCH] Allow seeing include $(x) when there is an include_top
 comment

Previously, seeing anything other than an interpolate variable
was not allowed.

Bug: 226974242
Test: go test
Change-Id: I48a060f9a3fd19cd67a114d2cb0756ab2be25ce1
---
 mk2rbc/mk2rbc.go      | 55 +++++++++++++++++++++++++++----------------
 mk2rbc/mk2rbc_test.go | 20 ++++++++++++++++
 mk2rbc/node.go        |  4 ++--
 3 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/mk2rbc/mk2rbc.go b/mk2rbc/mk2rbc.go
index 8f4fea4b1..57fc83fd1 100644
--- a/mk2rbc/mk2rbc.go
+++ b/mk2rbc/mk2rbc.go
@@ -811,35 +811,40 @@ func (ctx *parseContext) handleSubConfig(
 	//       rblf.inherit(handle, _e[0], _e[1])
 	//
 	var matchingPaths []string
-	varPath, ok := pathExpr.(*interpolateExpr)
-	if !ok {
+	var needsWarning = false
+	if interpolate, ok := pathExpr.(*interpolateExpr); ok {
+		pathPattern := []string{interpolate.chunks[0]}
+		for _, chunk := range interpolate.chunks[1:] {
+			if chunk != "" {
+				pathPattern = append(pathPattern, chunk)
+			}
+		}
+		if pathPattern[0] == "" && len(ctx.includeTops) > 0 {
+			// If pattern starts from the top. restrict it to the directories where
+			// we know inherit-product uses dynamically calculated path.
+			for _, p := range ctx.includeTops {
+				pathPattern[0] = p
+				matchingPaths = append(matchingPaths, ctx.findMatchingPaths(pathPattern)...)
+			}
+		} else {
+			matchingPaths = ctx.findMatchingPaths(pathPattern)
+		}
+		needsWarning = pathPattern[0] == "" && len(ctx.includeTops) == 0
+	} else if len(ctx.includeTops) > 0 {
+		for _, p := range ctx.includeTops {
+			matchingPaths = append(matchingPaths, ctx.findMatchingPaths([]string{p, ""})...)
+		}
+	} else {
 		return []starlarkNode{ctx.newBadNode(v, "inherit-product/include argument is too complex")}
 	}
 
-	pathPattern := []string{varPath.chunks[0]}
-	for _, chunk := range varPath.chunks[1:] {
-		if chunk != "" {
-			pathPattern = append(pathPattern, chunk)
-		}
-	}
-	if pathPattern[0] == "" && len(ctx.includeTops) > 0 {
-		// If pattern starts from the top. restrict it to the directories where
-		// we know inherit-product uses dynamically calculated path.
-		for _, p := range ctx.includeTops {
-			pathPattern[0] = p
-			matchingPaths = append(matchingPaths, ctx.findMatchingPaths(pathPattern)...)
-		}
-	} else {
-		matchingPaths = ctx.findMatchingPaths(pathPattern)
-	}
 	// Safeguard against $(call inherit-product,$(PRODUCT_PATH))
 	const maxMatchingFiles = 150
 	if len(matchingPaths) > maxMatchingFiles {
 		return []starlarkNode{ctx.newBadNode(v, "there are >%d files matching the pattern, please rewrite it", maxMatchingFiles)}
 	}
 
-	needsWarning := pathPattern[0] == "" && len(ctx.includeTops) == 0
-	res := inheritedDynamicModule{*varPath, []*moduleInfo{}, loadAlways, ctx.errorLocation(v), needsWarning}
+	res := inheritedDynamicModule{pathExpr, []*moduleInfo{}, loadAlways, ctx.errorLocation(v), needsWarning}
 	for _, p := range matchingPaths {
 		// A product configuration files discovered dynamically may attempt to inherit
 		// from another one which does not exist in this source tree. Prevent load errors
@@ -1587,6 +1592,16 @@ func transformNode(node starlarkNode, transformer func(expr starlarkExpr) starla
 		for _, n := range a.actions {
 			transformNode(n, transformer)
 		}
+	case *inheritNode:
+		if b, ok := a.module.(inheritedDynamicModule); ok {
+			b.path = b.path.transform(transformer)
+			a.module = b
+		}
+	case *includeNode:
+		if b, ok := a.module.(inheritedDynamicModule); ok {
+			b.path = b.path.transform(transformer)
+			a.module = b
+		}
 	}
 }
 
diff --git a/mk2rbc/mk2rbc_test.go b/mk2rbc/mk2rbc_test.go
index de7512927..b9a580df7 100644
--- a/mk2rbc/mk2rbc_test.go
+++ b/mk2rbc/mk2rbc_test.go
@@ -1142,6 +1142,11 @@ def init(g, handle):
 MY_PATH:=foo
 #RBC# include_top vendor/foo1
 $(call inherit-product,$(MY_PATH)/cfg.mk)
+#RBC# include_top vendor/foo1
+$(call inherit-product,$(MY_OTHER_PATH))
+#RBC# include_top vendor/foo1
+$(foreach f,$(MY_MAKEFILES), \
+	$(call inherit-product,$(f)))
 `,
 		expected: `load("//build/make/core:product_config.rbc", "rblf")
 load("//vendor/foo1:cfg.star|init", _cfg_init = "init")
@@ -1156,6 +1161,21 @@ def init(g, handle):
   if not _varmod_init:
     rblf.mkerror("product.mk", "Cannot find %s" % ("%s/cfg.mk" % g["MY_PATH"]))
   rblf.inherit(handle, _varmod, _varmod_init)
+  _entry = {
+    "vendor/foo1/cfg.mk": ("vendor/foo1/cfg", _cfg_init),
+  }.get(g.get("MY_OTHER_PATH", ""))
+  (_varmod, _varmod_init) = _entry if _entry else (None, None)
+  if not _varmod_init:
+    rblf.mkerror("product.mk", "Cannot find %s" % (g.get("MY_OTHER_PATH", "")))
+  rblf.inherit(handle, _varmod, _varmod_init)
+  for f in rblf.words(g.get("MY_MAKEFILES", "")):
+    _entry = {
+      "vendor/foo1/cfg.mk": ("vendor/foo1/cfg", _cfg_init),
+    }.get(f)
+    (_varmod, _varmod_init) = _entry if _entry else (None, None)
+    if not _varmod_init:
+      rblf.mkerror("product.mk", "Cannot find %s" % (f))
+    rblf.inherit(handle, _varmod, _varmod_init)
 `,
 	},
 	{
diff --git a/mk2rbc/node.go b/mk2rbc/node.go
index 7c39b9ead..a01abd8ac 100644
--- a/mk2rbc/node.go
+++ b/mk2rbc/node.go
@@ -83,7 +83,7 @@ func (im inheritedStaticModule) needsLoadCheck() bool {
 }
 
 type inheritedDynamicModule struct {
-	path             interpolateExpr
+	path             starlarkExpr
 	candidateModules []*moduleInfo
 	loadAlways       bool
 	location         ErrorLocation
@@ -120,7 +120,7 @@ func (i inheritedDynamicModule) emitSelect(gctx *generationContext) {
 }
 
 func (i inheritedDynamicModule) pathExpr() starlarkExpr {
-	return &i.path
+	return i.path
 }
 
 func (i inheritedDynamicModule) needsLoadCheck() bool {