Allow importing starlark code in makefiles

Adds a new `$(call run-starlark,my/starlark/file.bzl)` function that
will run the starlark file and set all the variables in the
variables_to_export_to_make dictionary as make variables.

Fixes: 280685526
Test: m nothing repeatedly causes no ninja regeneration, but touching all_versions.bzl does. go test, ./out/rbcrun -mode=rbc ./build/make/tests/run.rbc
Change-Id: Ic72e18dd28dba8233ba2dfb658b5d03ccece1bfd
diff --git a/tools/rbcrun/Android.bp b/tools/rbcrun/Android.bp
index fcc33ef..4fab858 100644
--- a/tools/rbcrun/Android.bp
+++ b/tools/rbcrun/Android.bp
@@ -34,6 +34,7 @@
     pkgPath: "rbcrun",
     deps: [
         "go-starlark-starlark",
+        "go-starlark-starlarkjson",
         "go-starlark-starlarkstruct",
         "go-starlark-starlarktest",
     ],
diff --git a/tools/rbcrun/host.go b/tools/rbcrun/host.go
index f2fda4e..a0fb9e1 100644
--- a/tools/rbcrun/host.go
+++ b/tools/rbcrun/host.go
@@ -24,13 +24,19 @@
 	"strings"
 
 	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkjson"
 	"go.starlark.net/starlarkstruct"
 )
 
-const callerDirKey = "callerDir"
+type ExecutionMode int
+const (
+	ExecutionModeRbc ExecutionMode = iota
+	ExecutionModeMake ExecutionMode = iota
+)
 
-var LoadPathRoot = "."
-var shellPath string
+const callerDirKey = "callerDir"
+const shellKey = "shell"
+const executionModeKey = "executionMode"
 
 type modentry struct {
 	globals starlark.StringDict
@@ -39,20 +45,66 @@
 
 var moduleCache = make(map[string]*modentry)
 
-var builtins starlark.StringDict
+var rbcBuiltins starlark.StringDict = starlark.StringDict{
+	"struct":   starlark.NewBuiltin("struct", starlarkstruct.Make),
+	// To convert find-copy-subdir and product-copy-files-by pattern
+	"rblf_find_files": starlark.NewBuiltin("rblf_find_files", find),
+	// To convert makefile's $(shell cmd)
+	"rblf_shell": starlark.NewBuiltin("rblf_shell", shell),
+	// Output to stderr
+	"rblf_log": starlark.NewBuiltin("rblf_log", log),
+	// To convert makefile's $(wildcard foo*)
+	"rblf_wildcard": starlark.NewBuiltin("rblf_wildcard", wildcard),
+}
 
-func moduleName2AbsPath(moduleName string, callerDir string) (string, error) {
-	path := moduleName
-	if ix := strings.LastIndex(path, ":"); ix >= 0 {
-		path = path[0:ix] + string(os.PathSeparator) + path[ix+1:]
+var makeBuiltins starlark.StringDict = starlark.StringDict{
+	"struct":   starlark.NewBuiltin("struct", starlarkstruct.Make),
+	"json": starlarkjson.Module,
+}
+
+// Takes a module name (the first argument to the load() function) and returns the path
+// it's trying to load, stripping out leading //, and handling leading :s.
+func cleanModuleName(moduleName string, callerDir string) (string, error) {
+	if strings.Count(moduleName, ":") > 1 {
+		return "", fmt.Errorf("at most 1 colon must be present in starlark path: %s", moduleName)
 	}
-	if strings.HasPrefix(path, "//") {
-		return filepath.Abs(filepath.Join(LoadPathRoot, path[2:]))
+
+	// We don't have full support for external repositories, but at least support skylib's dicts.
+	if moduleName == "@bazel_skylib//lib:dicts.bzl" {
+		return "external/bazel-skylib/lib/dicts.bzl", nil
+	}
+
+	localLoad := false
+	if strings.HasPrefix(moduleName, "@//") {
+		moduleName = moduleName[3:]
+	} else if strings.HasPrefix(moduleName, "//") {
+		moduleName = moduleName[2:]
 	} else if strings.HasPrefix(moduleName, ":") {
-		return filepath.Abs(filepath.Join(callerDir, path[1:]))
+		moduleName = moduleName[1:]
+		localLoad = true
 	} else {
-		return filepath.Abs(path)
+		return "", fmt.Errorf("load path must start with // or :")
 	}
+
+	if ix := strings.LastIndex(moduleName, ":"); ix >= 0 {
+		moduleName = moduleName[:ix] + string(os.PathSeparator) + moduleName[ix+1:]
+	}
+
+	if filepath.Clean(moduleName) != moduleName {
+		return "", fmt.Errorf("load path must be clean, found: %s, expected: %s", moduleName, filepath.Clean(moduleName))
+	}
+	if strings.HasPrefix(moduleName, "../") {
+		return "", fmt.Errorf("load path must not start with ../: %s", moduleName)
+	}
+	if strings.HasPrefix(moduleName, "/") {
+		return "", fmt.Errorf("load path starts with /, use // for a absolute path: %s", moduleName)
+	}
+
+	if localLoad {
+		return filepath.Join(callerDir, moduleName), nil
+	}
+
+	return moduleName, nil
 }
 
 // loader implements load statement. The format of the loaded module URI is
@@ -61,14 +113,18 @@
 // The presence of `|symbol` indicates that the loader should return a single 'symbol'
 // bound to None if file is missing.
 func loader(thread *starlark.Thread, module string) (starlark.StringDict, error) {
-	pipePos := strings.LastIndex(module, "|")
-	mustLoad := pipePos < 0
+	mode := thread.Local(executionModeKey).(ExecutionMode)
 	var defaultSymbol string
-	if !mustLoad {
-		defaultSymbol = module[pipePos+1:]
-		module = module[:pipePos]
+	mustLoad := true
+	if mode == ExecutionModeRbc {
+		pipePos := strings.LastIndex(module, "|")
+		mustLoad = pipePos < 0
+		if !mustLoad {
+			defaultSymbol = module[pipePos+1:]
+			module = module[:pipePos]
+		}
 	}
-	modulePath, err := moduleName2AbsPath(module, thread.Local(callerDirKey).(string))
+	modulePath, err := cleanModuleName(module, thread.Local(callerDirKey).(string))
 	if err != nil {
 		return nil, err
 	}
@@ -100,8 +156,17 @@
 			}
 
 			childThread.SetLocal(callerDirKey, filepath.Dir(modulePath))
-			globals, err := starlark.ExecFile(childThread, modulePath, nil, builtins)
-			e = &modentry{globals, err}
+			childThread.SetLocal(shellKey, thread.Local(shellKey))
+			childThread.SetLocal(executionModeKey, mode)
+			if mode == ExecutionModeRbc {
+				globals, err := starlark.ExecFile(childThread, modulePath, nil, rbcBuiltins)
+				e = &modentry{globals, err}
+			} else if mode == ExecutionModeMake {
+				globals, err := starlark.ExecFile(childThread, modulePath, nil, makeBuiltins)
+				e = &modentry{globals, err}
+			} else {
+				return nil, fmt.Errorf("unknown executionMode %d", mode)
+			}
 		} else {
 			e = &modentry{starlark.StringDict{defaultSymbol: starlark.None}, nil}
 		}
@@ -189,12 +254,13 @@
 // its output the same way as Make's $(shell ) function. The end-of-lines
 // ("\n" or "\r\n") are replaced with " " in the result, and the trailing
 // end-of-line is removed.
-func shell(_ *starlark.Thread, b *starlark.Builtin, args starlark.Tuple,
+func shell(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple,
 	kwargs []starlark.Tuple) (starlark.Value, error) {
 	var command string
 	if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &command); err != nil {
 		return starlark.None, err
 	}
+	shellPath := thread.Local(shellKey).(string)
 	if shellPath == "" {
 		return starlark.None,
 			fmt.Errorf("cannot run shell, /bin/sh is missing (running on Windows?)")
@@ -245,45 +311,68 @@
 	return starlark.None, nil
 }
 
-func setup() {
-	// Create the symbols that aid makefile conversion. See README.md
-	builtins = starlark.StringDict{
-		"struct":   starlark.NewBuiltin("struct", starlarkstruct.Make),
-		// To convert find-copy-subdir and product-copy-files-by pattern
-		"rblf_find_files": starlark.NewBuiltin("rblf_find_files", find),
-		// To convert makefile's $(shell cmd)
-		"rblf_shell": starlark.NewBuiltin("rblf_shell", shell),
-		// Output to stderr
-		"rblf_log": starlark.NewBuiltin("rblf_log", log),
-		// To convert makefile's $(wildcard foo*)
-		"rblf_wildcard": starlark.NewBuiltin("rblf_wildcard", wildcard),
-	}
-
-	// NOTE(asmundak): OS-specific. Behave similar to Linux `system` call,
-	// which always uses /bin/sh to run the command
-	shellPath = "/bin/sh"
-	if _, err := os.Stat(shellPath); err != nil {
-		shellPath = ""
-	}
-}
-
 // Parses, resolves, and executes a Starlark file.
 // filename and src parameters are as for starlark.ExecFile:
 // * filename is the name of the file to execute,
 //   and the name that appears in error messages;
 // * src is an optional source of bytes to use instead of filename
 //   (it can be a string, or a byte array, or an io.Reader instance)
-func Run(filename string, src interface{}) error {
-	setup()
+// Returns the top-level starlark variables, the list of starlark files loaded, and an error
+func Run(filename string, src interface{}, mode ExecutionMode) (starlark.StringDict, []string, error) {
+	// NOTE(asmundak): OS-specific. Behave similar to Linux `system` call,
+	// which always uses /bin/sh to run the command
+	shellPath := "/bin/sh"
+	if _, err := os.Stat(shellPath); err != nil {
+		shellPath = ""
+	}
+
 	mainThread := &starlark.Thread{
 		Name:  "main",
-		Print: func(_ *starlark.Thread, msg string) { fmt.Println(msg) },
+		Print: func(_ *starlark.Thread, msg string) {
+			if mode == ExecutionModeRbc {
+				// In rbc mode, rblf_log is used to print to stderr
+				fmt.Println(msg)
+			} else if mode == ExecutionModeMake {
+				fmt.Fprintln(os.Stderr, msg)
+			}
+		},
 		Load:  loader,
 	}
-	absPath, err := filepath.Abs(filename)
-	if err == nil {
-		mainThread.SetLocal(callerDirKey, filepath.Dir(absPath))
-		_, err = starlark.ExecFile(mainThread, absPath, src, builtins)
+	filename, err := filepath.Abs(filename)
+	if err != nil {
+		return nil, nil, err
 	}
-	return err
+	if wd, err := os.Getwd(); err == nil {
+		filename, err = filepath.Rel(wd, filename)
+		if err != nil {
+			return nil, nil, err
+		}
+		if strings.HasPrefix(filename, "../") {
+			return nil, nil, fmt.Errorf("path could not be made relative to workspace root: %s", filename)
+		}
+	} else {
+		return nil, nil, err
+	}
+
+	// Add top-level file to cache for cycle detection purposes
+	moduleCache[filename] = nil
+
+	var results starlark.StringDict
+	mainThread.SetLocal(callerDirKey, filepath.Dir(filename))
+	mainThread.SetLocal(shellKey, shellPath)
+	mainThread.SetLocal(executionModeKey, mode)
+	if mode == ExecutionModeRbc {
+		results, err = starlark.ExecFile(mainThread, filename, src, rbcBuiltins)
+	} else if mode == ExecutionModeMake {
+		results, err = starlark.ExecFile(mainThread, filename, src, makeBuiltins)
+	} else {
+		return results, nil, fmt.Errorf("unknown executionMode %d", mode)
+	}
+	loadedStarlarkFiles := make([]string, 0, len(moduleCache))
+	for file := range moduleCache {
+		loadedStarlarkFiles = append(loadedStarlarkFiles, file)
+	}
+	sort.Strings(loadedStarlarkFiles)
+
+	return results, loadedStarlarkFiles, err
 }
diff --git a/tools/rbcrun/host_test.go b/tools/rbcrun/host_test.go
index e109c02..10cac62 100644
--- a/tools/rbcrun/host_test.go
+++ b/tools/rbcrun/host_test.go
@@ -54,7 +54,6 @@
 
 // Common setup for the tests: create thread, change to the test directory
 func testSetup(t *testing.T) *starlark.Thread {
-	setup()
 	thread := &starlark.Thread{
 		Load: func(thread *starlark.Thread, module string) (starlark.StringDict, error) {
 			if module == "assert.star" {
@@ -78,7 +77,6 @@
 	// In order to use "assert.star" from go/starlark.net/starlarktest in the tests, provide:
 	//  * load function that handles "assert.star"
 	//  * starlarktest.DataFile function that finds its location
-	setup()
 	if err := os.Chdir(dataDir()); err != nil {
 		t.Fatal(err)
 	}
@@ -92,7 +90,9 @@
 	starlarktest.SetReporter(thread, t)
 	_, thisSrcFile, _, _ := runtime.Caller(0)
 	filename := filepath.Join(filepath.Dir(thisSrcFile), starFile)
-	if _, err := starlark.ExecFile(thread, filename, nil, builtins); err != nil {
+	thread.SetLocal(executionModeKey, ExecutionModeRbc)
+	thread.SetLocal(shellKey, "/bin/sh")
+	if _, err := starlark.ExecFile(thread, filename, nil, rbcBuiltins); err != nil {
 		if err, ok := err.(*starlark.EvalError); ok {
 			t.Fatal(err.Backtrace())
 		}
@@ -103,7 +103,7 @@
 func TestFileOps(t *testing.T) {
 	// TODO(asmundak): convert this to use exerciseStarlarkTestFile
 	thread := testSetup(t)
-	if _, err := starlark.ExecFile(thread, "file_ops.star", nil, builtins); err != nil {
+	if _, err := starlark.ExecFile(thread, "file_ops.star", nil, rbcBuiltins); err != nil {
 		if err, ok := err.(*starlark.EvalError); ok {
 			t.Fatal(err.Backtrace())
 		}
@@ -122,9 +122,12 @@
 		}
 	}
 	dir := dataDir()
+	if err := os.Chdir(filepath.Dir(dir)); err != nil {
+		t.Fatal(err)
+	}
 	thread.SetLocal(callerDirKey, dir)
-	LoadPathRoot = filepath.Dir(dir)
-	if _, err := starlark.ExecFile(thread, "load.star", nil, builtins); err != nil {
+	thread.SetLocal(executionModeKey, ExecutionModeRbc)
+	if _, err := starlark.ExecFile(thread, "testdata/load.star", nil, rbcBuiltins); err != nil {
 		if err, ok := err.(*starlark.EvalError); ok {
 			t.Fatal(err.Backtrace())
 		}
diff --git a/tools/rbcrun/rbcrun/rbcrun.go b/tools/rbcrun/rbcrun/rbcrun.go
index 8dd0f46..b5182f0 100644
--- a/tools/rbcrun/rbcrun/rbcrun.go
+++ b/tools/rbcrun/rbcrun/rbcrun.go
@@ -17,18 +17,22 @@
 import (
 	"flag"
 	"fmt"
-	"go.starlark.net/starlark"
 	"os"
 	"rbcrun"
+	"regexp"
+	"strings"
+
+	"go.starlark.net/starlark"
 )
 
 var (
+	modeFlag  = flag.String("mode", "", "the general behavior of rbcrun. Can be \"rbc\" or \"make\". Required.")
 	rootdir  = flag.String("d", ".", "the value of // for load paths")
 	perfFile = flag.String("perf", "", "save performance data")
+	identifierRe = regexp.MustCompile("[a-zA-Z_][a-zA-Z0-9_]*")
 )
 
-func main() {
-	flag.Parse()
+func getEntrypointStarlarkFile() string {
 	filename := ""
 
 	for _, arg := range flag.Args() {
@@ -42,8 +46,108 @@
 		flag.Usage()
 		os.Exit(1)
 	}
-	if stat, err := os.Stat(*rootdir); os.IsNotExist(err) || !stat.IsDir() {
-		quit("%s is not a directory\n", *rootdir)
+	return filename
+}
+
+func getMode() rbcrun.ExecutionMode {
+	switch *modeFlag {
+	case "rbc":
+		return rbcrun.ExecutionModeRbc
+	case "make":
+		return rbcrun.ExecutionModeMake
+	case "":
+		quit("-mode flag is required.")
+	default:
+		quit("Unknown -mode value %q, expected 1 of \"rbc\", \"make\"", *modeFlag)
+	}
+	return rbcrun.ExecutionModeMake
+}
+
+var makeStringReplacer = strings.NewReplacer("#", "\\#", "$", "$$")
+
+func cleanStringForMake(s string) (string, error) {
+	if strings.ContainsAny(s, "\\\n") {
+		// \\ in make is literally \\, not a single \, so we can't allow them.
+		// \<newline> in make will produce a space, not a newline.
+		return "", fmt.Errorf("starlark strings exported to make cannot contain backslashes or newlines")
+	}
+	return makeStringReplacer.Replace(s), nil
+}
+
+func getValueInMakeFormat(value starlark.Value, allowLists bool) (string, error) {
+	switch v := value.(type) {
+	case starlark.String:
+		if cleanedValue, err := cleanStringForMake(v.GoString()); err == nil {
+			return cleanedValue, nil
+		} else {
+			return "", err
+		}
+	case starlark.Int:
+		return v.String(), nil
+	case *starlark.List:
+		if !allowLists {
+			return "", fmt.Errorf("nested lists are not allowed to be exported from starlark to make, flatten the list in starlark first")
+		}
+		result := ""
+		for i := 0; i < v.Len(); i++ {
+			value, err := getValueInMakeFormat(v.Index(i), false)
+			if err != nil {
+				return "", err
+			}
+			if i > 0 {
+				result += " "
+			}
+			result += value
+		}
+		return result, nil
+	default:
+		return "", fmt.Errorf("only starlark strings, ints, and lists of strings/ints can be exported to make. Please convert all other types in starlark first. Found type: %s", value.Type())
+	}
+}
+
+func printVarsInMakeFormat(globals starlark.StringDict) error {
+	// We could just directly export top level variables by name instead of going through
+	// a variables_to_export_to_make dictionary, but that wouldn't allow for exporting a
+	// runtime-defined number of variables to make. This can be important because dictionaries
+	// in make are often represented by a unique variable for every key in the dictionary.
+	variablesValue, ok := globals["variables_to_export_to_make"]
+	if !ok {
+		return fmt.Errorf("expected top-level starlark file to have a \"variables_to_export_to_make\" variable")
+	}
+	variables, ok := variablesValue.(*starlark.Dict)
+	if !ok {
+		return fmt.Errorf("expected variables_to_export_to_make to be a dict, got %s", variablesValue.Type())
+	}
+
+	for _, varTuple := range variables.Items() {
+		varNameStarlark, ok := varTuple.Index(0).(starlark.String)
+		if !ok {
+			return fmt.Errorf("all keys in variables_to_export_to_make must be strings, but got %q", varTuple.Index(0).Type())
+		}
+		varName := varNameStarlark.GoString()
+		if !identifierRe.MatchString(varName) {
+			return fmt.Errorf("all variables at the top level starlark file must be valid c identifiers, but got %q", varName)
+		}
+		if varName == "LOADED_STARLARK_FILES" {
+			return fmt.Errorf("the name LOADED_STARLARK_FILES is reserved for use by the starlark interpreter")
+		}
+		valueMake, err := getValueInMakeFormat(varTuple.Index(1), true)
+		if err != nil {
+			return err
+		}
+		// The :=$= is special Kati syntax that means "set and make readonly"
+		fmt.Printf("%s :=$= %s\n", varName, valueMake)
+	}
+	return nil
+}
+
+func main() {
+	flag.Parse()
+	filename := getEntrypointStarlarkFile()
+	mode := getMode()
+
+	if os.Chdir(*rootdir) != nil {
+		quit("could not chdir to %s\n", *rootdir)
 	}
 	if *perfFile != "" {
 		pprof, err := os.Create(*perfFile)
@@ -55,8 +159,7 @@
 			quit("%s\n", err)
 		}
 	}
-	rbcrun.LoadPathRoot = *rootdir
-	err := rbcrun.Run(filename, nil)
+	variables, loadedStarlarkFiles, err := rbcrun.Run(filename, nil, mode)
 	rc := 0
 	if *perfFile != "" {
 		if err2 := starlark.StopProfile(); err2 != nil {
@@ -71,6 +174,12 @@
 			quit("%s\n", err)
 		}
 	}
+	if mode == rbcrun.ExecutionModeMake {
+		if err := printVarsInMakeFormat(variables); err != nil {
+			quit("%s\n", err)
+		}
+		fmt.Printf("LOADED_STARLARK_FILES := %s\n", strings.Join(loadedStarlarkFiles, " "))
+	}
 	os.Exit(rc)
 }
 
diff --git a/tools/rbcrun/testdata/module1.star b/tools/rbcrun/testdata/module1.star
index be04f75..02919a0 100644
--- a/tools/rbcrun/testdata/module1.star
+++ b/tools/rbcrun/testdata/module1.star
@@ -2,6 +2,6 @@
 load("assert.star", "assert")
 
 # Make sure that builtins are defined for the loaded module, too
-assert.true(rblf_wildcard("module1.star"))
-assert.true(not rblf_wildcard("no_such file"))
+assert.true(rblf_wildcard("testdata/module1.star"))
+assert.true(not rblf_wildcard("testdata/no_such file"))
 test = "module1"