Add map of aliases into import declaration, update tests.

onflow · Dec 16, 2024 · 739a8a9 · 739a8a9
1 parent 027432c
commit 739a8a9
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 15 deletions.
diff --git a/ast/import.go b/ast/import.go
@@ -31,6 +31,7 @@ import (
 type ImportDeclaration struct {
 	Location    common.Location
 	Identifiers []Identifier
+	Aliases     map[string]string
 	Range
 	LocationPos Position
 }
@@ -41,6 +42,7 @@ var _ Declaration = &ImportDeclaration{}
 func NewImportDeclaration(
 	gauge common.MemoryGauge,
 	identifiers []Identifier,
+	aliases map[string]string,
 	location common.Location,
 	declRange Range,
 	locationPos Position,
@@ -49,6 +51,7 @@ func NewImportDeclaration(
 
 	return &ImportDeclaration{
 		Identifiers: identifiers,
+		Aliases:     aliases,
 		Location:    location,
 		Range:       declRange,
 		LocationPos: locationPos,

diff --git a/parser/declaration.go b/parser/declaration.go
@@ -652,6 +652,7 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 	startPosition := p.current.StartPos
 
 	var identifiers []ast.Identifier
+	aliases := make(map[string]string)
 
 	var location common.Location
 	var locationPos ast.Position
@@ -684,7 +685,8 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 		endPos = identifier.EndPosition(p.memoryGauge)
 	}
 
-	parseOptionalImportAlias := func() error {
+	// pass in the identifier which would be aliased
+	parseOptionalImportAlias := func(identifier ast.Identifier) error {
 		// stop early if the current token is not as
 		if string(p.currentTokenSource()) != KeywordAs {
 			return nil
@@ -695,7 +697,7 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 		switch p.current.Type {
 		case lexer.TokenIdentifier:
 			identifierAlias := p.tokenToIdentifier(p.current)
-			identifiers = append(identifiers, identifierAlias)
+			aliases[identifier.Identifier] = identifierAlias.Identifier
 
 			// Skip the alias
 			p.nextSemanticToken()
@@ -729,15 +731,10 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 	}
 
 	parseMoreIdentifiers := func() error {
-		expectCommaOrFrom := false
+		expectCommaOrFrom := true
 
 		atEnd := false
 		for !atEnd {
-			p.nextSemanticToken()
-
-			// Parse optional alias
-			parseOptionalImportAlias()
-
 			switch p.current.Type {
 			case lexer.TokenComma:
 				if !expectCommaOrFrom {
@@ -749,6 +746,7 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 					)
 				}
 				expectCommaOrFrom = false
+				p.nextSemanticToken()
 
 			case lexer.TokenIdentifier:
 
@@ -782,6 +780,10 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 
 				identifier := p.tokenToIdentifier(p.current)
 				identifiers = append(identifiers, identifier)
+				p.nextSemanticToken()
+
+				// Parse optional alias
+				parseOptionalImportAlias(identifier)
 
 				expectCommaOrFrom = true
 
@@ -841,7 +843,7 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 		// Skip the identifier
 		p.nextSemanticToken()
 		// Parse optional alias
-		parseOptionalImportAlias()
+		parseOptionalImportAlias(identifier)
 
 		switch p.current.Type {
 		case lexer.TokenComma:
@@ -883,6 +885,7 @@ func parseImportDeclaration(p *parser) (*ast.ImportDeclaration, error) {
 	return ast.NewImportDeclaration(
 		p.memoryGauge,
 		identifiers,
+		aliases,
 		location,
 		ast.NewRange(
 			p.memoryGauge,

diff --git a/parser/declaration_test.go b/parser/declaration_test.go
@@ -2436,36 +2436,123 @@ func TestParseImportDeclaration(t *testing.T) {
 
 		t.Parallel()
 
-		_, errs := testParseDeclarations(`
+		result, errs := testParseDeclarations(`
 			import foo as bar, lorem from 0x42
 		`)
 		require.Empty(t, errs)
 
-		// assert result after deciding how to parse
+		AssertEqualWithDiff(t,
+			[]ast.Declaration{
+				&ast.ImportDeclaration{
+					Identifiers: []ast.Identifier{
+						{
+							Identifier: "foo",
+							Pos:        ast.Position{Line: 2, Column: 10, Offset: 11},
+						},
+						{
+							Identifier: "lorem",
+							Pos:        ast.Position{Line: 2, Column: 22, Offset: 23},
+						},
+					},
+					Aliases: map[string]string{
+						"foo": "bar",
+					},
+					Location: common.AddressLocation{
+						Address: common.MustBytesToAddress([]byte{0x42}),
+					},
+					LocationPos: ast.Position{Line: 2, Column: 33, Offset: 34},
+					Range: ast.Range{
+						StartPos: ast.Position{Line: 2, Column: 3, Offset: 4},
+						EndPos:   ast.Position{Line: 2, Column: 36, Offset: 37},
+					},
+				},
+			},
+			result,
+		)
 	})
 
 	t.Run("multiple import alias", func(t *testing.T) {
 
 		t.Parallel()
 
-		_, errs := testParseDeclarations(`
+		result, errs := testParseDeclarations(`
 			import foo as bar, lorem as ipsum from 0x42
 		`)
 		require.Empty(t, errs)
 
-		// assert result after deciding how to parse
+		AssertEqualWithDiff(t,
+			[]ast.Declaration{
+				&ast.ImportDeclaration{
+					Identifiers: []ast.Identifier{
+						{
+							Identifier: "foo",
+							Pos:        ast.Position{Line: 2, Column: 10, Offset: 11},
+						},
+						{
+							Identifier: "lorem",
+							Pos:        ast.Position{Line: 2, Column: 22, Offset: 23},
+						},
+					},
+					Aliases: map[string]string{
+						"foo":   "bar",
+						"lorem": "ipsum",
+					},
+					Location: common.AddressLocation{
+						Address: common.MustBytesToAddress([]byte{0x42}),
+					},
+					LocationPos: ast.Position{Line: 2, Column: 42, Offset: 43},
+					Range: ast.Range{
+						StartPos: ast.Position{Line: 2, Column: 3, Offset: 4},
+						EndPos:   ast.Position{Line: 2, Column: 45, Offset: 46},
+					},
+				},
+			},
+			result,
+		)
 	})
 
 	t.Run("combination import aliases", func(t *testing.T) {
 
 		t.Parallel()
 
-		_, errs := testParseDeclarations(`
+		result, errs := testParseDeclarations(`
 			import foo as bar, test as from, from from 0x42
 		`)
 		require.Empty(t, errs)
 
-		// assert result after deciding how to parse
+		AssertEqualWithDiff(t,
+			[]ast.Declaration{
+				&ast.ImportDeclaration{
+					Identifiers: []ast.Identifier{
+						{
+							Identifier: "foo",
+							Pos:        ast.Position{Line: 2, Column: 10, Offset: 11},
+						},
+						{
+							Identifier: "test",
+							Pos:        ast.Position{Line: 2, Column: 22, Offset: 23},
+						},
+						{
+							Identifier: "from",
+							Pos:        ast.Position{Line: 2, Column: 36, Offset: 37},
+						},
+					},
+					Aliases: map[string]string{
+						"foo":  "bar",
+						"test": "from",
+					},
+					Location: common.AddressLocation{
+						Address: common.MustBytesToAddress([]byte{0x42}),
+					},
+					LocationPos: ast.Position{Line: 2, Column: 46, Offset: 47},
+					Range: ast.Range{
+						StartPos: ast.Position{Line: 2, Column: 3, Offset: 4},
+						EndPos:   ast.Position{Line: 2, Column: 49, Offset: 50},
+					},
+				},
+			},
+			result,
+		)
 	})
 }