starlark_import.go 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306
  1. // Copyright 2023 Google Inc. All rights reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package starlark_import
  15. import (
  16. "fmt"
  17. "os"
  18. "path/filepath"
  19. "sort"
  20. "strings"
  21. "sync"
  22. "time"
  23. "go.starlark.net/starlark"
  24. "go.starlark.net/starlarkjson"
  25. "go.starlark.net/starlarkstruct"
  26. )
  27. func init() {
  28. go func() {
  29. startTime := time.Now()
  30. v, d, err := runStarlarkFile("//build/bazel/constants_exported_to_soong.bzl")
  31. endTime := time.Now()
  32. //fmt.Fprintf(os.Stderr, "starlark run time: %s\n", endTime.Sub(startTime).String())
  33. globalResult.Set(starlarkResult{
  34. values: v,
  35. ninjaDeps: d,
  36. err: err,
  37. startTime: startTime,
  38. endTime: endTime,
  39. })
  40. }()
  41. }
  42. type starlarkResult struct {
  43. values starlark.StringDict
  44. ninjaDeps []string
  45. err error
  46. startTime time.Time
  47. endTime time.Time
  48. }
  49. // setOnce wraps a value and exposes Set() and Get() accessors for it.
  50. // The Get() calls will block until a Set() has been called.
  51. // A second call to Set() will panic.
  52. // setOnce must be created using newSetOnce()
  53. type setOnce[T any] struct {
  54. value T
  55. lock sync.Mutex
  56. wg sync.WaitGroup
  57. isSet bool
  58. }
  59. func (o *setOnce[T]) Set(value T) {
  60. o.lock.Lock()
  61. defer o.lock.Unlock()
  62. if o.isSet {
  63. panic("Value already set")
  64. }
  65. o.value = value
  66. o.isSet = true
  67. o.wg.Done()
  68. }
  69. func (o *setOnce[T]) Get() T {
  70. if !o.isSet {
  71. o.wg.Wait()
  72. }
  73. return o.value
  74. }
  75. func newSetOnce[T any]() *setOnce[T] {
  76. result := &setOnce[T]{}
  77. result.wg.Add(1)
  78. return result
  79. }
  80. var globalResult = newSetOnce[starlarkResult]()
  81. func GetStarlarkValue[T any](key string) (T, error) {
  82. result := globalResult.Get()
  83. if result.err != nil {
  84. var zero T
  85. return zero, result.err
  86. }
  87. if !result.values.Has(key) {
  88. var zero T
  89. return zero, fmt.Errorf("a starlark variable by that name wasn't found, did you update //build/bazel/constants_exported_to_soong.bzl?")
  90. }
  91. return Unmarshal[T](result.values[key])
  92. }
  93. func GetNinjaDeps() ([]string, error) {
  94. result := globalResult.Get()
  95. if result.err != nil {
  96. return nil, result.err
  97. }
  98. return result.ninjaDeps, nil
  99. }
  100. func getTopDir() (string, error) {
  101. // It's hard to communicate the top dir to this package in any other way than reading the
  102. // arguments directly, because we need to know this at package initialization time. Many
  103. // soong constants that we'd like to read from starlark are initialized during package
  104. // initialization.
  105. for i, arg := range os.Args {
  106. if arg == "--top" {
  107. if i < len(os.Args)-1 && os.Args[i+1] != "" {
  108. return os.Args[i+1], nil
  109. }
  110. }
  111. }
  112. // When running tests, --top is not passed. Instead, search for the top dir manually
  113. cwd, err := os.Getwd()
  114. if err != nil {
  115. return "", err
  116. }
  117. for cwd != "/" {
  118. if _, err := os.Stat(filepath.Join(cwd, "build/soong/soong_ui.bash")); err == nil {
  119. return cwd, nil
  120. }
  121. cwd = filepath.Dir(cwd)
  122. }
  123. return "", fmt.Errorf("could not find top dir")
  124. }
  125. const callerDirKey = "callerDir"
  126. type modentry struct {
  127. globals starlark.StringDict
  128. err error
  129. }
  130. func unsupportedMethod(t *starlark.Thread, fn *starlark.Builtin, _ starlark.Tuple, _ []starlark.Tuple) (starlark.Value, error) {
  131. return nil, fmt.Errorf("%sthis file is read by soong, and must therefore be pure starlark and include only constant information. %q is not allowed", t.CallStack().String(), fn.Name())
  132. }
  133. var builtins = starlark.StringDict{
  134. "aspect": starlark.NewBuiltin("aspect", unsupportedMethod),
  135. "glob": starlark.NewBuiltin("glob", unsupportedMethod),
  136. "json": starlarkjson.Module,
  137. "provider": starlark.NewBuiltin("provider", unsupportedMethod),
  138. "rule": starlark.NewBuiltin("rule", unsupportedMethod),
  139. "struct": starlark.NewBuiltin("struct", starlarkstruct.Make),
  140. "select": starlark.NewBuiltin("select", unsupportedMethod),
  141. "transition": starlark.NewBuiltin("transition", unsupportedMethod),
  142. }
  143. // Takes a module name (the first argument to the load() function) and returns the path
  144. // it's trying to load, stripping out leading //, and handling leading :s.
  145. func cleanModuleName(moduleName string, callerDir string) (string, error) {
  146. if strings.Count(moduleName, ":") > 1 {
  147. return "", fmt.Errorf("at most 1 colon must be present in starlark path: %s", moduleName)
  148. }
  149. // We don't have full support for external repositories, but at least support skylib's dicts.
  150. if moduleName == "@bazel_skylib//lib:dicts.bzl" {
  151. return "external/bazel-skylib/lib/dicts.bzl", nil
  152. }
  153. localLoad := false
  154. if strings.HasPrefix(moduleName, "@//") {
  155. moduleName = moduleName[3:]
  156. } else if strings.HasPrefix(moduleName, "//") {
  157. moduleName = moduleName[2:]
  158. } else if strings.HasPrefix(moduleName, ":") {
  159. moduleName = moduleName[1:]
  160. localLoad = true
  161. } else {
  162. return "", fmt.Errorf("load path must start with // or :")
  163. }
  164. if ix := strings.LastIndex(moduleName, ":"); ix >= 0 {
  165. moduleName = moduleName[:ix] + string(os.PathSeparator) + moduleName[ix+1:]
  166. }
  167. if filepath.Clean(moduleName) != moduleName {
  168. return "", fmt.Errorf("load path must be clean, found: %s, expected: %s", moduleName, filepath.Clean(moduleName))
  169. }
  170. if strings.HasPrefix(moduleName, "../") {
  171. return "", fmt.Errorf("load path must not start with ../: %s", moduleName)
  172. }
  173. if strings.HasPrefix(moduleName, "/") {
  174. return "", fmt.Errorf("load path starts with /, use // for a absolute path: %s", moduleName)
  175. }
  176. if localLoad {
  177. return filepath.Join(callerDir, moduleName), nil
  178. }
  179. return moduleName, nil
  180. }
  181. // loader implements load statement. The format of the loaded module URI is
  182. //
  183. // [//path]:base
  184. //
  185. // The file path is $ROOT/path/base if path is present, <caller_dir>/base otherwise.
  186. func loader(thread *starlark.Thread, module string, topDir string, moduleCache map[string]*modentry, moduleCacheLock *sync.Mutex, filesystem map[string]string) (starlark.StringDict, error) {
  187. modulePath, err := cleanModuleName(module, thread.Local(callerDirKey).(string))
  188. if err != nil {
  189. return nil, err
  190. }
  191. moduleCacheLock.Lock()
  192. e, ok := moduleCache[modulePath]
  193. if e == nil {
  194. if ok {
  195. moduleCacheLock.Unlock()
  196. return nil, fmt.Errorf("cycle in load graph")
  197. }
  198. // Add a placeholder to indicate "load in progress".
  199. moduleCache[modulePath] = nil
  200. moduleCacheLock.Unlock()
  201. childThread := &starlark.Thread{Name: "exec " + module, Load: thread.Load}
  202. // Cheating for the sake of testing:
  203. // propagate starlarktest's Reporter key, otherwise testing
  204. // the load function may cause panic in starlarktest code.
  205. const testReporterKey = "Reporter"
  206. if v := thread.Local(testReporterKey); v != nil {
  207. childThread.SetLocal(testReporterKey, v)
  208. }
  209. childThread.SetLocal(callerDirKey, filepath.Dir(modulePath))
  210. if filesystem != nil {
  211. globals, err := starlark.ExecFile(childThread, filepath.Join(topDir, modulePath), filesystem[modulePath], builtins)
  212. e = &modentry{globals, err}
  213. } else {
  214. globals, err := starlark.ExecFile(childThread, filepath.Join(topDir, modulePath), nil, builtins)
  215. e = &modentry{globals, err}
  216. }
  217. // Update the cache.
  218. moduleCacheLock.Lock()
  219. moduleCache[modulePath] = e
  220. }
  221. moduleCacheLock.Unlock()
  222. return e.globals, e.err
  223. }
  224. // Run runs the given starlark file and returns its global variables and a list of all starlark
  225. // files that were loaded. The top dir for starlark's // is found via getTopDir().
  226. func runStarlarkFile(filename string) (starlark.StringDict, []string, error) {
  227. topDir, err := getTopDir()
  228. if err != nil {
  229. return nil, nil, err
  230. }
  231. return runStarlarkFileWithFilesystem(filename, topDir, nil)
  232. }
  233. func runStarlarkFileWithFilesystem(filename string, topDir string, filesystem map[string]string) (starlark.StringDict, []string, error) {
  234. if !strings.HasPrefix(filename, "//") && !strings.HasPrefix(filename, ":") {
  235. filename = "//" + filename
  236. }
  237. filename, err := cleanModuleName(filename, "")
  238. if err != nil {
  239. return nil, nil, err
  240. }
  241. moduleCache := make(map[string]*modentry)
  242. moduleCache[filename] = nil
  243. moduleCacheLock := &sync.Mutex{}
  244. mainThread := &starlark.Thread{
  245. Name: "main",
  246. Print: func(_ *starlark.Thread, msg string) {
  247. // Ignore prints
  248. },
  249. Load: func(thread *starlark.Thread, module string) (starlark.StringDict, error) {
  250. return loader(thread, module, topDir, moduleCache, moduleCacheLock, filesystem)
  251. },
  252. }
  253. mainThread.SetLocal(callerDirKey, filepath.Dir(filename))
  254. var result starlark.StringDict
  255. if filesystem != nil {
  256. result, err = starlark.ExecFile(mainThread, filepath.Join(topDir, filename), filesystem[filename], builtins)
  257. } else {
  258. result, err = starlark.ExecFile(mainThread, filepath.Join(topDir, filename), nil, builtins)
  259. }
  260. return result, sortedStringKeys(moduleCache), err
  261. }
  262. func sortedStringKeys(m map[string]*modentry) []string {
  263. s := make([]string, 0, len(m))
  264. for k := range m {
  265. s = append(s, k)
  266. }
  267. sort.Strings(s)
  268. return s
  269. }