Record occurrences of nested types and imported members

sema/check_member_expression.go

@@ -306,6 +306,7 @@ func (checker *Checker) visitMember(expression *ast.MemberExpression, isAssignme
 
 	if checker.PositionInfo != nil {
 		checker.PositionInfo.recordMemberOccurrence(
+			checker.memoryGauge,
 			accessedType,
 			identifier,
 			identifierStartPosition,

sema/checker.go

@@ -1272,6 +1272,10 @@ func (checker *Checker) convertNominalType(t *ast.NominalType) Type {
 			)
 			return InvalidType
 		}
+
+		if checker.PositionInfo != nil && identifier.Identifier != "" {
+			checker.recordNestedTypeReferenceOccurrence(identifier, ty)
+		}
 	}
 
 	return ty
@@ -1428,6 +1432,15 @@ func (checker *Checker) recordVariableReferenceOccurrence(startPos, endPos ast.P
 	)
 }
 
+func (checker *Checker) recordNestedTypeReferenceOccurrence(identifier ast.Identifier, nestedType Type) {
+	checker.PositionInfo.recordNestedTypeReferenceOccurrence(
+		checker.memoryGauge,
+		checker.Elaboration,
+		identifier,
+		nestedType,
+	)
+}
+
 func (checker *Checker) recordVariableDeclarationOccurrence(name string, variable *Variable) {
 	checker.PositionInfo.recordVariableDeclarationOccurrence(
 		checker.memoryGauge,

sema/elaboration.go

@@ -323,12 +323,11 @@ func (e *Elaboration) SetCompositeDeclarationType(
 	e.compositeDeclarationTypes[declaration] = compositeType
 }
 
-func (e *Elaboration) CompositeTypeDeclaration(compositeType *CompositeType) (decl ast.CompositeLikeDeclaration, ok bool) {
+func (e *Elaboration) CompositeTypeDeclaration(compositeType *CompositeType) ast.CompositeLikeDeclaration {
 	if e.compositeTypeDeclarations == nil {
-		return
+		return nil
 	}
-	decl, ok = e.compositeTypeDeclarations[compositeType]
-	return
+	return e.compositeTypeDeclarations[compositeType]
 }
 
 func (e *Elaboration) SetCompositeTypeDeclaration(
@@ -419,6 +418,34 @@ func (e *Elaboration) EntitlementMapTypeDeclaration(entitlementMapType *Entitlem
 	return decl
 }
 
+func (e *Elaboration) DeclarationForType(ty Type) ast.Declaration {
+	// Each case checks the concrete return value against nil before returning,
+	// rather than `return e.XTypeDeclaration(t)`, to avoid Go's typed nil behavior
+	switch t := ty.(type) {
+	case *CompositeType:
+		if decl := e.CompositeTypeDeclaration(t); decl != nil {
+			return decl
+		}
+
+	case *InterfaceType:
+		if decl := e.InterfaceTypeDeclaration(t); decl != nil {
+			return decl
+		}
+
+	case *EntitlementType:
+		if decl := e.EntitlementTypeDeclaration(t); decl != nil {
+			return decl
+		}
+
+	case *EntitlementMapType:
+		if decl := e.EntitlementMapTypeDeclaration(t); decl != nil {
+			return decl
+		}
+	}
+
+	return nil
+}
+
 func (e *Elaboration) ConstructorFunctionType(initializer *ast.SpecialFunctionDeclaration) *FunctionType {
 	if e.constructorFunctionTypes == nil {
 		return nil

sema/elaboration_test.go

@@ -0,0 +1,123 @@
+/*
+ * Cadence - The resource-oriented smart contract programming language
+ *
+ * Copyright Flow Foundation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package sema_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/onflow/cadence/sema"
+	. "github.com/onflow/cadence/test_utils/sema_utils"
+)
+
+func TestElaborationDeclarationForType(t *testing.T) {
+
+	t.Parallel()
+
+	checker, err := ParseAndCheck(t, `
+        struct S {}
+
+        struct interface I {}
+
+        entitlement E
+
+        entitlement mapping M {}
+    `)
+	require.NoError(t, err)
+
+	elaboration := checker.Elaboration
+
+	compositeType := RequireGlobalType(t, elaboration, "S").(*sema.CompositeType)
+	interfaceType := RequireGlobalType(t, elaboration, "I").(*sema.InterfaceType)
+	entitlementType := elaboration.EntitlementType("S.test.E")
+	require.NotNil(t, entitlementType)
+	entitlementMapType := elaboration.EntitlementMapType("S.test.M")
+	require.NotNil(t, entitlementMapType)
+
+	t.Run("composite type", func(t *testing.T) {
+		t.Parallel()
+
+		expected := checker.Program.CompositeDeclarations()[0]
+		require.Same(t, expected, elaboration.DeclarationForType(compositeType))
+	})
+
+	t.Run("interface type", func(t *testing.T) {
+		t.Parallel()
+
+		expected := checker.Program.InterfaceDeclarations()[0]
+		require.Same(t, expected, elaboration.DeclarationForType(interfaceType))
+	})
+
+	t.Run("entitlement type", func(t *testing.T) {
+		t.Parallel()
+
+		expected := checker.Program.EntitlementDeclarations()[0]
+		require.Same(t, expected, elaboration.DeclarationForType(entitlementType))
+	})
+
+	t.Run("entitlement map type", func(t *testing.T) {
+		t.Parallel()
+
+		expected := checker.Program.EntitlementMappingDeclarations()[0]
+		require.Same(t, expected, elaboration.DeclarationForType(entitlementMapType))
+	})
+
+	t.Run("unhandled type kind", func(t *testing.T) {
+		t.Parallel()
+
+		// Anything outside the four handled kinds should fall through the switch
+		// and return an untyped nil ast.Declaration.
+		result := elaboration.DeclarationForType(sema.IntType)
+		require.True(t, result == nil)
+	})
+
+	// An empty elaboration has no declaration recorded for these types.
+	// The helpers return a typed-nil value, but DeclarationForType must
+	// return an untyped nil ast.Declaration so callers can compare against nil.
+	emptyElaboration := sema.NewElaboration(nil)
+
+	t.Run("composite type not in elaboration", func(t *testing.T) {
+		t.Parallel()
+
+		result := emptyElaboration.DeclarationForType(compositeType)
+		require.True(t, result == nil)
+	})
+
+	t.Run("interface type not in elaboration", func(t *testing.T) {
+		t.Parallel()
+
+		result := emptyElaboration.DeclarationForType(interfaceType)
+		require.True(t, result == nil)
+	})
+
+	t.Run("entitlement type not in elaboration", func(t *testing.T) {
+		t.Parallel()
+
+		result := emptyElaboration.DeclarationForType(entitlementType)
+		require.True(t, result == nil)
+	})
+
+	t.Run("entitlement map type not in elaboration", func(t *testing.T) {
+		t.Parallel()
+
+		result := emptyElaboration.DeclarationForType(entitlementMapType)
+		require.True(t, result == nil)
+	})
+}

sema/occurrences_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/onflow/cadence/ast"
 	"github.com/onflow/cadence/common"
 	"github.com/onflow/cadence/sema"
 	. "github.com/onflow/cadence/test_utils/sema_utils"
@@ -359,3 +360,85 @@ nextMatcher:
 		assert.NotNil(t, checker.PositionInfo.Occurrences.Find(matcher.EndPos))
 	}
 }
+
+func TestCheckNestedTypeOccurrenceSameFile(t *testing.T) {
+
+	t.Parallel()
+
+	checker, err := ParseAndCheckWithOptions(t, `
+        access(all) contract A {
+            access(all) struct S {}
+            access(all) fun b(_ s: A.S) {}
+        }
+        `,
+		ParseAndCheckOptions{
+			CheckerConfig: &sema.Config{
+				PositionInfoEnabled: true,
+			},
+		},
+	)
+	require.NoError(t, err)
+
+	contractA := RequireGlobalType(t, checker.Elaboration, "A").(*sema.CompositeType)
+	sType, _ := contractA.GetNestedTypes().Get("S")
+
+	occurrence := checker.PositionInfo.Occurrences.Find(sema.Position{Line: 4, Column: 37})
+	require.NotNil(t, occurrence)
+
+	assert.Equal(t,
+		&sema.Origin{
+			Type:            sType,
+			StartPos:        &ast.Position{Offset: 65, Line: 3, Column: 31},
+			EndPos:          &ast.Position{Offset: 65, Line: 3, Column: 31},
+			DeclarationKind: common.DeclarationKindStructure,
+			Occurrences: []ast.Range{
+				{
+					StartPos: ast.Position{Offset: 107, Line: 4, Column: 37},
+					EndPos:   ast.Position{Offset: 107, Line: 4, Column: 37},
+				},
+			},
+		},
+		occurrence.Origin,
+	)
+}
+
+func TestCheckNestedInterfaceTypeOccurrence(t *testing.T) {
+
+	t.Parallel()
+
+	checker, err := ParseAndCheckWithOptions(t, `
+        access(all) contract A {
+            access(all) struct interface SI {}
+            access(all) fun b(_ s: {A.SI}) {}
+        }
+        `,
+		ParseAndCheckOptions{
+			CheckerConfig: &sema.Config{
+				PositionInfoEnabled: true,
+			},
+		},
+	)
+	require.NoError(t, err)
+
+	contractA := RequireGlobalType(t, checker.Elaboration, "A").(*sema.CompositeType)
+	siType, _ := contractA.GetNestedTypes().Get("SI")
+
+	occurrence := checker.PositionInfo.Occurrences.Find(sema.Position{Line: 4, Column: 38})
+	require.NotNil(t, occurrence)
+
+	assert.Equal(t,
+		&sema.Origin{
+			Type:            siType,
+			StartPos:        &ast.Position{Offset: 75, Line: 3, Column: 41},
+			EndPos:          &ast.Position{Offset: 76, Line: 3, Column: 42},
+			DeclarationKind: common.DeclarationKindStructureInterface,
+			Occurrences: []ast.Range{
+				{
+					StartPos: ast.Position{Offset: 119, Line: 4, Column: 38},
+					EndPos:   ast.Position{Offset: 120, Line: 4, Column: 39},
+				},
+			},
+		},
+		occurrence.Origin,
+	)
+}