From 9672d86a8772216311a1bc4dede23c41664c8bc7 Mon Sep 17 00:00:00 2001
From: Colin Cross <ccross@android.com>
Date: Mon, 30 Dec 2019 18:38:20 -0800
Subject: [PATCH] Add scoped module factories

Add RegisterScopedModuleType to LoadHookContext that registers
a module type factory that will be visible for the remainder of
the file.  Also add ModuleFactories that returns the globally
register module factories.

Test: all blueprint tests
Change-Id: If646a73befe3b8e45c4b0984531c6a39ddc8d066
---
 context.go    | 11 ++++++++---
 module_ctx.go | 39 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/context.go b/context.go
index 0e4f586..5d10aec 100644
--- a/context.go
+++ b/context.go
@@ -670,12 +670,14 @@ func (c *Context) ParseFileList(rootDir string, filePaths []string,
 
 		addedCh := make(chan struct{})
 
+		var scopedModuleFactories map[string]ModuleFactory
+
 		var addModule func(module *moduleInfo) []error
 		addModule = func(module *moduleInfo) (errs []error) {
 			moduleCh <- newModuleInfo{module, addedCh}
 			<-addedCh
 			var newModules []*moduleInfo
-			newModules, errs = runAndRemoveLoadHooks(c, config, module)
+			newModules, errs = runAndRemoveLoadHooks(c, config, module, &scopedModuleFactories)
 			if len(errs) > 0 {
 				return errs
 			}
@@ -691,7 +693,7 @@ func (c *Context) ParseFileList(rootDir string, filePaths []string,
 		for _, def := range file.Defs {
 			switch def := def.(type) {
 			case *parser.Module:
-				module, errs := c.processModuleDef(def, file.Name)
+				module, errs := c.processModuleDef(def, file.Name, scopedModuleFactories)
 				if len(errs) == 0 && module != nil {
 					errs = addModule(module)
 				}
@@ -1338,9 +1340,12 @@ func (c *Context) newModule(factory ModuleFactory) *moduleInfo {
 }
 
 func (c *Context) processModuleDef(moduleDef *parser.Module,
-	relBlueprintsFile string) (*moduleInfo, []error) {
+	relBlueprintsFile string, scopedModuleFactories map[string]ModuleFactory) (*moduleInfo, []error) {
 
 	factory, ok := c.moduleFactories[moduleDef.Type]
+	if !ok && scopedModuleFactories != nil {
+		factory, ok = scopedModuleFactories[moduleDef.Type]
+	}
 	if !ok {
 		if c.ignoreUnknownModuleTypes {
 			return nil, nil
diff --git a/module_ctx.go b/module_ctx.go
index 8b52b50..4a365f5 100644
--- a/module_ctx.go
+++ b/module_ctx.go
@@ -179,6 +179,9 @@ type EarlyModuleContext interface {
 	// Namespace returns the Namespace object provided by the NameInterface set by Context.SetNameInterface, or the
 	// default SimpleNameInterface if Context.SetNameInterface was not called.
 	Namespace() Namespace
+
+	// ModuleFactories returns a map of all of the global ModuleFactories by name.
+	ModuleFactories() map[string]ModuleFactory
 }
 
 type BaseModuleContext interface {
@@ -602,6 +605,14 @@ func (m *baseModuleContext) AddNinjaFileDeps(deps ...string) {
 	m.ninjaFileDeps = append(m.ninjaFileDeps, deps...)
 }
 
+func (m *baseModuleContext) ModuleFactories() map[string]ModuleFactory {
+	ret := make(map[string]ModuleFactory)
+	for k, v := range m.context.moduleFactories {
+		ret[k] = v
+	}
+	return ret
+}
+
 func (m *moduleContext) ModuleSubDir() string {
 	return m.module.variantName
 }
@@ -1009,6 +1020,10 @@ type LoadHookContext interface {
 	// CreateModule creates a new module by calling the factory method for the specified moduleType, and applies
 	// the specified property structs to it as if the properties were set in a blueprint file.
 	CreateModule(ModuleFactory, ...interface{}) Module
+
+	// RegisterScopedModuleType creates a new module type that is scoped to the current Blueprints
+	// file.
+	RegisterScopedModuleType(name string, factory ModuleFactory)
 }
 
 func (l *loadHookContext) CreateModule(factory ModuleFactory, props ...interface{}) Module {
@@ -1031,9 +1046,26 @@ func (l *loadHookContext) CreateModule(factory ModuleFactory, props ...interface
 	return module.logicModule
 }
 
+func (l *loadHookContext) RegisterScopedModuleType(name string, factory ModuleFactory) {
+	if _, exists := l.context.moduleFactories[name]; exists {
+		panic(fmt.Errorf("A global module type named %q already exists", name))
+	}
+
+	if _, exists := (*l.scopedModuleFactories)[name]; exists {
+		panic(fmt.Errorf("A module type named %q already exists in this scope", name))
+	}
+
+	if *l.scopedModuleFactories == nil {
+		(*l.scopedModuleFactories) = make(map[string]ModuleFactory)
+	}
+
+	(*l.scopedModuleFactories)[name] = factory
+}
+
 type loadHookContext struct {
 	baseModuleContext
-	newModules []*moduleInfo
+	newModules            []*moduleInfo
+	scopedModuleFactories *map[string]ModuleFactory
 }
 
 type LoadHook func(ctx LoadHookContext)
@@ -1057,8 +1089,8 @@ func AddLoadHook(module Module, hook LoadHook) {
 	*hooks = append(*hooks, hook)
 }
 
-func runAndRemoveLoadHooks(ctx *Context, config interface{},
-	module *moduleInfo) (newModules []*moduleInfo, errs []error) {
+func runAndRemoveLoadHooks(ctx *Context, config interface{}, module *moduleInfo,
+	scopedModuleFactories *map[string]ModuleFactory) (newModules []*moduleInfo, errs []error) {
 
 	if v, exists := pendingHooks.Load(module.logicModule); exists {
 		hooks := v.(*[]LoadHook)
@@ -1068,6 +1100,7 @@ func runAndRemoveLoadHooks(ctx *Context, config interface{},
 				config:  config,
 				module:  module,
 			},
+			scopedModuleFactories: scopedModuleFactories,
 		}
 
 		for _, hook := range *hooks {