Merge #11575

11575: Simplify `findFunctionSchema` r=iwahbe a=iwahbe

Instead of indirecting through .NET function names, just use the schema token to lookup the function.

Co-authored-by: Ian Wahbe <ian@wahbe.com>

pkg/codegen/dotnet/gen_program.go

@@ -260,23 +260,23 @@ func (g *generator) warnf(location *hcl.Range, reason string, args ...interface{
 	})
 }
 
-func (g *generator) findFunctionSchema(function string, location *hcl.Range) (*schema.Function, bool) {
-	function = LowerCamelCase(function)
+func (g *generator) findFunctionSchema(token string, location *hcl.Range) (*schema.Function, bool) {
 	for _, pkg := range g.program.PackageReferences() {
-		for it := pkg.Functions().Range(); it.Next(); {
-			if strings.HasSuffix(it.Token(), function) {
-				fn, err := it.Function()
-
-				if err != nil {
-					g.warnf(location, "Could not find function schema for '%s'; err %s", function, err.Error())
-					return nil, false
-				}
-
-				return fn, true
-			}
+		fn, ok, err := pcl.LookupFunction(pkg, token)
+		if !ok {
+			continue
 		}
+		if err != nil {
+			g.diagnostics = append(g.diagnostics, &hcl.Diagnostic{
+				Severity: hcl.DiagWarning,
+				Summary:  fmt.Sprintf("Could not find function schema for '%s'", token),
+				Detail:   err.Error(),
+				Subject:  location,
+			})
+			return nil, false
+		}
+		return fn, true
 	}
-
 	return nil, false
 }
 
@@ -287,12 +287,8 @@ func (g *generator) isFunctionInvoke(localVariable *pcl.LocalVariable) (*schema.
 		call := value.(*model.FunctionCallExpression)
 		switch call.Name {
 		case pcl.Invoke:
-			args := call.Args[0]
-			_, fullFunctionName := g.functionName(args)
-			functionNameParts := strings.Split(fullFunctionName, ".")
-			functionName := functionNameParts[len(functionNameParts)-1]
-			location := value.SyntaxNode().Range().Ptr()
-			return g.findFunctionSchema(functionName, location)
+			token := call.Args[0].(*model.TemplateExpression).Parts[0].(*model.LiteralValueExpression).Value.AsString()
+			return g.findFunctionSchema(token, call.Args[0].SyntaxNode().Range().Ptr())
 		}
 	}
 

pkg/codegen/pcl/binder_schema.go

@@ -33,6 +33,8 @@ type packageSchema struct {
 	schema schema.PackageReference
 
 	// These maps map from canonical tokens to actual tokens.
+	//
+	// Both maps take `nil` to mean uninitialized.
 	resourceTokenMap map[string]string
 	functionTokenMap map[string]string
 }
@@ -42,9 +44,25 @@ type packageOpts struct {
 	pluginDownloadURL string
 }
 
+// Lookup a PCL invoke token in a schema.
+func LookupFunction(pkg schema.PackageReference, token string) (*schema.Function, bool, error) {
+	s, _, ok, err := newPackageSchema(pkg).LookupFunction(token)
+	return s, ok, err
+}
+
+// Lookup a PCL resource token in a schema.
+func LookupResource(pkg schema.PackageReference, token string) (*schema.Resource, bool, error) {
+	r, _, ok, err := newPackageSchema(pkg).LookupResource(token)
+	return r, ok, err
+}
+
 func (ps *packageSchema) LookupFunction(token string) (*schema.Function, string, bool, error) {
 	contract.Assert(ps != nil)
 
+	if ps.functionTokenMap == nil {
+		ps.initFunctionMap()
+	}
+
 	schemaToken, ok := ps.functionTokenMap[token]
 	if !ok {
 		token = canonicalizeToken(token, ps.schema)
@@ -61,6 +79,10 @@ func (ps *packageSchema) LookupFunction(token string) (*schema.Function, string,
 func (ps *packageSchema) LookupResource(token string) (*schema.Resource, string, bool, error) {
 	contract.Assert(ps != nil)
 
+	if ps.resourceTokenMap == nil {
+		ps.initResourceMap()
+	}
+
 	schemaToken, ok := ps.resourceTokenMap[token]
 	if !ok {
 		token = canonicalizeToken(token, ps.schema)
@@ -74,6 +96,26 @@ func (ps *packageSchema) LookupResource(token string) (*schema.Resource, string,
 	return res, token, ok, err
 }
 
+func (ps *packageSchema) initFunctionMap() {
+	functionTokenMap := map[string]string{}
+	for it := ps.schema.Functions().Range(); it.Next(); {
+		functionTokenMap[canonicalizeToken(it.Token(), ps.schema)] = it.Token()
+	}
+	ps.functionTokenMap = functionTokenMap
+}
+
+func (ps *packageSchema) initResourceMap() {
+	resourceTokenMap := map[string]string{}
+	for it := ps.schema.Resources().Range(); it.Next(); {
+		resourceTokenMap[canonicalizeToken(it.Token(), ps.schema)] = it.Token()
+	}
+	ps.resourceTokenMap = resourceTokenMap
+}
+
+func newPackageSchema(pkg schema.PackageReference) *packageSchema {
+	return &packageSchema{schema: pkg}
+}
+
 type PackageInfo struct {
 	name    string
 	version string
@@ -122,20 +164,7 @@ func (c *PackageCache) loadPackageSchema(loader schema.Loader, name, version str
 		return nil, err
 	}
 
-	resourceTokenMap := map[string]string{}
-	for it := pkg.Resources().Range(); it.Next(); {
-		resourceTokenMap[canonicalizeToken(it.Token(), pkg)] = it.Token()
-	}
-	functionTokenMap := map[string]string{}
-	for it := pkg.Functions().Range(); it.Next(); {
-		functionTokenMap[canonicalizeToken(it.Token(), pkg)] = it.Token()
-	}
-
-	schema := &packageSchema{
-		schema:           pkg,
-		resourceTokenMap: resourceTokenMap,
-		functionTokenMap: functionTokenMap,
-	}
+	schema := newPackageSchema(pkg)
 
 	c.m.Lock()
 	defer c.m.Unlock()