AST VariableDeclaration Tuple and Return Statement Fix + Resolver Imp…

…rovements (#108)

* More reference resolution, return statement fix
* Type description void if return does not have expression
* Wth...
* VariableDeclaration tuple were missing...
* Fixing builder_test

ast/builder_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -27,6 +28,7 @@ func TestAstBuilderFromSourceAsString(t *testing.T) {
 
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
+
 			parser, err := solgo.NewParserFromSources(context.TODO(), testCase.sources)
 			assert.NoError(t, err)
 			assert.NotNil(t, parser)
@@ -55,6 +57,7 @@ func TestAstBuilderFromSourceAsString(t *testing.T) {
 			assert.Equal(t, errsExpected, errs)
 			assert.Equal(t, int(testCase.unresolvedReferences), astBuilder.GetResolver().GetUnprocessedCount())
 			assert.Equal(t, len(astBuilder.GetResolver().GetUnprocessedNodes()), astBuilder.GetResolver().GetUnprocessedCount())
+
 			for _, sourceUnit := range astBuilder.GetRoot().GetSourceUnits() {
 				prettyJson, err := utils.ToJSONPretty(sourceUnit)
 				assert.NoError(t, err)
@@ -126,11 +129,18 @@ func TestAstBuilderFromSourceAsString(t *testing.T) {
 }
 
 func recursiveTest(t *testing.T, node Node[NodeType]) {
+	if node == nil {
+		return
+	}
+
 	assert.NotNil(t, node.GetNodes(), fmt.Sprintf("Node %T has nil nodes", node))
 	assert.GreaterOrEqual(t, node.GetId(), int64(0), fmt.Sprintf("Node %T has empty id", node))
 	assert.NotNil(t, node.GetType(), fmt.Sprintf("Node %T has empty type", node))
 	assert.NotNil(t, node.GetSrc(), fmt.Sprintf("Node %T has empty GetSrc()", node))
-	assert.NotNil(t, node.GetTypeDescription(), fmt.Sprintf("Node %T has not defined GetTypeDescription()", node))
+
+	if reflect.TypeOf(node).String() != "*ast.ReturnStatement" {
+		assert.NotNil(t, node.GetTypeDescription(), fmt.Sprintf("Node %T has not defined GetTypeDescription()", node))
+	}
 
 	if contract, ok := node.(*Contract); ok {
 		assert.GreaterOrEqual(t, len(contract.GetBaseContracts()), 0)
@@ -351,8 +361,10 @@ func recursiveReferenceDescriptorSetTest(t *testing.T, node Node[NodeType]) {
 	node.SetReferenceDescriptor(0, &TypeDescription{})
 
 	for _, childNode := range node.GetNodes() {
-		childNode.SetReferenceDescriptor(0, nil)
-		childNode.SetReferenceDescriptor(0, &TypeDescription{})
-		recursiveReferenceDescriptorSetTest(t, childNode)
+		if childNode != nil {
+			childNode.SetReferenceDescriptor(0, nil)
+			childNode.SetReferenceDescriptor(0, &TypeDescription{})
+			recursiveReferenceDescriptorSetTest(t, childNode)
+		}
 	}
 }

ast/index_access.go

@@ -119,7 +119,9 @@ func (i *IndexAccess) ToProto() NodeType {
 	}
 
 	for _, td := range i.GetTypeDescriptions() {
-		proto.TypeDescriptions = append(proto.TypeDescriptions, td.ToProto())
+		if td != nil {
+			proto.TypeDescriptions = append(proto.TypeDescriptions, td.ToProto())
+		}
 	}
 
 	return NewTypedStruct(&proto, "IndexAccess")
@@ -193,6 +195,14 @@ func (i *IndexAccess) buildTypeDescription() *TypeDescription {
 	typeIdentifiers := make([]string, 0)
 
 	for _, paramType := range i.GetTypeDescriptions() {
+		// REMOVE-LATER: It's a fix because sometimes forward-path is not quite working at this stage...
+		// For example, defining state variables at end of the contract instead of the top :explosion:
+		if paramType == nil {
+			typeStrings = append(typeStrings, "unknown")
+			typeIdentifiers = append(typeIdentifiers, "$_t_unknown")
+			continue
+		}
+
 		if strings.Contains(paramType.TypeString, "literal_string") {
 			typeStrings = append(typeStrings, "string memory")
 			typeIdentifiers = append(typeIdentifiers, "_"+paramType.TypeIdentifier)

ast/primary_expression.go

@@ -539,6 +539,13 @@ func (p *PrimaryExpression) Parse(
 		}
 	}
 
+	if p.TypeDescription == nil {
+		if refId, refTypeDescription := p.GetResolver().ResolveByNode(p, p.Name); refTypeDescription != nil {
+			p.ReferencedDeclaration = refId
+			p.TypeDescription = refTypeDescription
+		}
+	}
+
 	return p
 }
 

ast/reference.go

@@ -149,6 +149,14 @@ func (r *Resolver) Resolve() []error {
 						),
 					)
 				}
+			} else {
+				errors = append(
+					errors,
+					fmt.Errorf(
+						"unable to find node by id %d - name: %s - type: %v",
+						nodeId, node.Name, reflect.TypeOf(node.Node),
+					),
+				)
 			}
 		}
 
@@ -483,6 +491,14 @@ func (r *Resolver) byRecursiveSearch(node Node[NodeType], name string) (Node[Nod
 		if nodeCtx.GetName() == name {
 			return nodeCtx, nodeCtx.GetTypeDescription()
 		}
+	case *Assignment:
+		if nodeCtx.GetRightExpression() != nil {
+			if expr, ok := nodeCtx.GetRightExpression().(*PrimaryExpression); ok {
+				if expr.GetName() == name {
+					return expr, expr.GetTypeDescription()
+				}
+			}
+		}
 	}
 
 	for _, n := range node.GetNodes() {

ast/return.go

@@ -58,10 +58,14 @@ func (r *ReturnStatement) GetFunctionReturnParameters() int64 {
 
 // GetTypeDescription returns the type description of the ReturnStatement's expression.
 func (r *ReturnStatement) GetTypeDescription() *TypeDescription {
-	if r.Expression != nil {
-		return r.Expression.GetTypeDescription()
+	if r.Expression == nil {
+		return &TypeDescription{
+			TypeString:     "void",
+			TypeIdentifier: "$_t_return_void",
+		}
 	}
-	return nil
+
+	return r.Expression.GetTypeDescription()
 }
 
 // GetNodes returns a list of child nodes contained in the ReturnStatement.
@@ -76,8 +80,14 @@ func (r *ReturnStatement) ToProto() NodeType {
 		NodeType:                 r.GetType(),
 		Src:                      r.Src.ToProto(),
 		FunctionReturnParameters: r.GetFunctionReturnParameters(),
-		Expression:               r.GetExpression().ToProto().(*v3.TypedStruct),
-		TypeDescription:          r.GetTypeDescription().ToProto(),
+	}
+
+	if r.GetExpression() != nil {
+		proto.Expression = r.GetExpression().ToProto().(*v3.TypedStruct)
+	}
+
+	if r.GetTypeDescription() != nil {
+		proto.TypeDescription = r.GetTypeDescription().ToProto()
 	}
 
 	return NewTypedStruct(&proto, "Return")
@@ -106,7 +116,10 @@ func (r *ReturnStatement) Parse(
 		r.FunctionReturnParameters = fnCtx.GetId()
 	}
 
-	expression := NewExpression(r.ASTBuilder)
-	r.Expression = expression.Parse(unit, contractNode, fnNode, bodyNode, nil, nil, ctx.Expression())
+	if ctx.Expression() != nil {
+		expression := NewExpression(r.ASTBuilder)
+		r.Expression = expression.Parse(unit, contractNode, fnNode, bodyNode, nil, nil, ctx.Expression())
+	}
+
 	return r
 }

ast/sources_test.go

@@ -42,6 +42,24 @@ func getSourceTestCases(t *testing.T) []struct {
 			expectedProto:        tests.ReadJsonBytesForTest(t, "ast/Empty.solgo.ast.proto").Content,
 			unresolvedReferences: 0,
 		},
+		{
+			name:       "PAPA Contract",
+			outputPath: "contracts/papa/",
+			sources: &solgo.Sources{
+				SourceUnits: []*solgo.SourceUnit{
+					{
+						Name:    "ERC20",
+						Path:    tests.ReadContractFileForTest(t, "contracts/papa/Token").Path,
+						Content: tests.ReadContractFileForTest(t, "contracts/papa/Token").Content,
+					},
+				},
+				EntrySourceUnitName: "Token",
+				LocalSourcesPath:    buildFullPath("../sources/"),
+			},
+			expectedAst:          tests.ReadJsonBytesForTest(t, "contracts/papa/Token.solgo.ast").Content,
+			expectedProto:        tests.ReadJsonBytesForTest(t, "contracts/papa/Token.solgo.ast.proto").Content,
+			unresolvedReferences: 0,
+		},
 		{
 			name:       "Simple Storage Contract Test",
 			outputPath: "ast/",

ast/variable.go

@@ -136,6 +136,15 @@ func (v *VariableDeclaration) Parse(
 		v.Assignments = append(v.Assignments, declaration.GetId())
 	}
 
+	if ctx.VariableDeclarationTuple() != nil {
+		for _, declarationCtx := range ctx.VariableDeclarationTuple().AllVariableDeclaration() {
+			declaration := NewDeclaration(v.ASTBuilder)
+			declaration.ParseVariableDeclaration(unit, contractNode, fnNode, bodyNode, v, declarationCtx)
+			v.Declarations = append(v.Declarations, declaration)
+			v.Assignments = append(v.Assignments, declaration.GetId())
+		}
+	}
+
 	if ctx.Expression() != nil {
 		expression := NewExpression(v.ASTBuilder)
 		v.InitialValue = expression.Parse(unit, contractNode, fnNode, bodyNode, v, nil, ctx.Expression())

data/tests/abi/TransparentUpgradeableProxy.abi.json

@@ -1,5 +1,5 @@
 {
-	"entry_contract_id": 2309,
+	"entry_contract_id": 2337,
 	"entry_contract_name": "TransparentUpgradeableProxy",
 	"contracts_count": 13,
 	"contracts": {

data/tests/abi/TransparentUpgradeableProxy.abi.proto.json

@@ -1,5 +1,5 @@
 {
-	"entryContractId": 2309,
+	"entryContractId": 2337,
 	"entryContractName": "TransparentUpgradeableProxy",
 	"contractsCount": 13,
 	"contracts": {