Refactoring + Speed improvements (#36)

* Refactoring + Speed improvements
Removed unnecessary return values in node stack
Added an import cache bringing massive perfomance improvements

* Fix imports

* Rename some vars

pkg/nodestack/nodestack.go

@@ -25,19 +25,18 @@ func (s *NodeStack) Clone() *NodeStack {
 	}
 }
 
-func (s *NodeStack) Push(n ast.Node) *NodeStack {
+func (s *NodeStack) Push(n ast.Node) {
 	s.Stack = append(s.Stack, n)
-	return s
 }
 
-func (s *NodeStack) Pop() (*NodeStack, ast.Node) {
+func (s *NodeStack) Pop() ast.Node {
 	l := len(s.Stack)
 	if l == 0 {
-		return s, nil
+		return nil
 	}
 	n := s.Stack[l-1]
 	s.Stack = s.Stack[:l-1]
-	return s, n
+	return n
 }
 
 func (s *NodeStack) Peek() ast.Node {
@@ -54,14 +53,14 @@ func (s *NodeStack) IsEmpty() bool {
 func (s *NodeStack) BuildIndexList() []string {
 	var indexList []string
 	for !s.IsEmpty() {
-		_, curr := s.Pop()
+		curr := s.Pop()
 		switch curr := curr.(type) {
 		case *ast.SuperIndex:
-			s = s.Push(curr.Index)
+			s.Push(curr.Index)
 			indexList = append(indexList, "super")
 		case *ast.Index:
-			s = s.Push(curr.Index)
-			s = s.Push(curr.Target)
+			s.Push(curr.Index)
+			s.Push(curr.Target)
 		case *ast.LiteralString:
 			indexList = append(indexList, curr.Value)
 		case *ast.Self:

pkg/processing/find_bind.go

@@ -6,10 +6,8 @@ import (
 )
 
 func FindBindByIdViaStack(stack *nodestack.NodeStack, id ast.Identifier) *ast.LocalBind {
-	stack = stack.Clone()
-	for !stack.IsEmpty() {
-		_, curr := stack.Pop()
-		switch curr := curr.(type) {
+	for _, node := range stack.Stack {
+		switch curr := node.(type) {
 		case *ast.Local:
 			for _, bind := range curr.Binds {
 				if bind.Variable == id {
@@ -23,7 +21,6 @@ func FindBindByIdViaStack(stack *nodestack.NodeStack, id ast.Identifier) *ast.Lo
 				}
 			}
 		}
-
 	}
 	return nil
 }

pkg/processing/find_field.go

@@ -59,8 +59,7 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 	} else if start == "std" {
 		return nil, fmt.Errorf("cannot get definition of std lib")
 	} else if strings.Contains(start, ".") {
-		rootNode, _, _ := vm.ImportAST("", start)
-		foundDesugaredObjects = findTopLevelObjects(nodestack.NewNodeStack(rootNode), vm)
+		foundDesugaredObjects = findTopLevelObjectsInFile(vm, start, "")
 	} else if start == "$" {
 		sameFileOnly = true
 		foundDesugaredObjects = findTopLevelObjects(nodestack.NewNodeStack(stack.From), vm)
@@ -88,8 +87,7 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 			foundDesugaredObjects = findTopLevelObjects(tmpStack, vm)
 		case *ast.Import:
 			filename := bodyNode.File.Value
-			rootNode, _, _ := vm.ImportAST("", filename)
-			foundDesugaredObjects = findTopLevelObjects(nodestack.NewNodeStack(rootNode), vm)
+			foundDesugaredObjects = findTopLevelObjectsInFile(vm, filename, "")
 		case *ast.Index:
 			tempStack := nodestack.NewNodeStack(bodyNode)
 			indexList = append(tempStack.BuildIndexList(), indexList...)
@@ -103,7 +101,7 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 		index := indexList[0]
 		indexList = indexList[1:]
 		foundFields := findObjectFieldsInObjects(foundDesugaredObjects, index)
-		foundDesugaredObjects = foundDesugaredObjects[:0]
+		foundDesugaredObjects = nil
 		if len(foundFields) == 0 {
 			return nil, fmt.Errorf("field %s was not found in ast.DesugaredObject", index)
 		}
@@ -119,31 +117,9 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 			return ranges, nil
 		}
 
-		// Unpack:
-		// - Binary nodes. A field could be either in the left or right side of the binary
-		// - Self nodes. We want the object self refers to, not the self node itself
-		var fieldNodes []ast.Node
-		for _, foundField := range foundFields {
-			switch fieldNode := foundField.Body.(type) {
-			case *ast.Self:
-				filename := fieldNode.LocRange.FileName
-				rootNode, _, _ := vm.ImportAST("", filename)
-				tmpStack, err := FindNodeByPosition(rootNode, fieldNode.LocRange.Begin)
-				if err != nil {
-					return nil, err
-				}
-				for !tmpStack.IsEmpty() {
-					_, node := tmpStack.Pop()
-					if _, ok := node.(*ast.DesugaredObject); ok {
-						fieldNodes = append(fieldNodes, node)
-					}
-				}
-			case *ast.Binary:
-				fieldNodes = append(fieldNodes, fieldNode.Right)
-				fieldNodes = append(fieldNodes, fieldNode.Left)
-			default:
-				fieldNodes = append(fieldNodes, fieldNode)
-			}
+		fieldNodes, err := unpackFieldNodes(vm, foundFields)
+		if err != nil {
+			return nil, err
 		}
 
 		for _, fieldNode := range fieldNodes {
@@ -155,7 +131,7 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 				}
 				foundDesugaredObjects = append(foundDesugaredObjects, varReference.(*ast.DesugaredObject))
 			case *ast.DesugaredObject:
-				stack = stack.Push(fieldNode)
+				stack.Push(fieldNode)
 				foundDesugaredObjects = append(foundDesugaredObjects, findDesugaredObjectFromStack(stack))
 			case *ast.Index:
 				tempStack := nodestack.NewNodeStack(fieldNode)
@@ -168,15 +144,46 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 				return result, err
 			case *ast.Import:
 				filename := fieldNode.File.Value
-				rootNode, _, _ := vm.ImportAST(string(fieldNode.Loc().File.DiagnosticFileName), filename)
-				foundDesugaredObjects = append(foundDesugaredObjects, findTopLevelObjects(nodestack.NewNodeStack(rootNode), vm)...)
+				newObjs := findTopLevelObjectsInFile(vm, filename, string(fieldNode.Loc().File.DiagnosticFileName))
+				foundDesugaredObjects = append(foundDesugaredObjects, newObjs...)
 			}
 		}
 	}
 
 	return ranges, nil
 }
 
+// unpackFieldNodes extracts nodes from fields
+// - Binary nodes. A field could be either in the left or right side of the binary
+// - Self nodes. We want the object self refers to, not the self node itself
+func unpackFieldNodes(vm *jsonnet.VM, fields []*ast.DesugaredObjectField) ([]ast.Node, error) {
+	var fieldNodes []ast.Node
+	for _, foundField := range fields {
+		switch fieldNode := foundField.Body.(type) {
+		case *ast.Self:
+			filename := fieldNode.LocRange.FileName
+			rootNode, _, _ := vm.ImportAST("", filename)
+			tmpStack, err := FindNodeByPosition(rootNode, fieldNode.LocRange.Begin)
+			if err != nil {
+				return nil, err
+			}
+			for !tmpStack.IsEmpty() {
+				node := tmpStack.Pop()
+				if _, ok := node.(*ast.DesugaredObject); ok {
+					fieldNodes = append(fieldNodes, node)
+				}
+			}
+		case *ast.Binary:
+			fieldNodes = append(fieldNodes, fieldNode.Right)
+			fieldNodes = append(fieldNodes, fieldNode.Left)
+		default:
+			fieldNodes = append(fieldNodes, fieldNode)
+		}
+	}
+
+	return fieldNodes, nil
+}
+
 func findObjectFieldsInObjects(objectNodes []*ast.DesugaredObject, index string) []*ast.DesugaredObjectField {
 	var matchingFields []*ast.DesugaredObjectField
 	for _, object := range objectNodes {
@@ -207,7 +214,7 @@ func findObjectFieldInObject(objectNode *ast.DesugaredObject, index string) *ast
 
 func findDesugaredObjectFromStack(stack *nodestack.NodeStack) *ast.DesugaredObject {
 	for !stack.IsEmpty() {
-		_, curr := stack.Pop()
+		curr := stack.Pop()
 		switch curr := curr.(type) {
 		case *ast.DesugaredObject:
 			return curr
@@ -216,47 +223,6 @@ func findDesugaredObjectFromStack(stack *nodestack.NodeStack) *ast.DesugaredObje
 	return nil
 }
 
-// Find all ast.DesugaredObject's from NodeStack
-func findTopLevelObjects(stack *nodestack.NodeStack, vm *jsonnet.VM) []*ast.DesugaredObject {
-	var objects []*ast.DesugaredObject
-	for !stack.IsEmpty() {
-		_, curr := stack.Pop()
-		switch curr := curr.(type) {
-		case *ast.DesugaredObject:
-			objects = append(objects, curr)
-		case *ast.Binary:
-			stack = stack.Push(curr.Left)
-			stack = stack.Push(curr.Right)
-		case *ast.Local:
-			stack = stack.Push(curr.Body)
-		case *ast.Import:
-			filename := curr.File.Value
-			rootNode, _, _ := vm.ImportAST(string(curr.Loc().File.DiagnosticFileName), filename)
-			stack = stack.Push(rootNode)
-		case *ast.Index:
-			container := stack.Peek()
-			if containerObj, containerIsObj := container.(*ast.DesugaredObject); containerIsObj {
-				indexValue, indexIsString := curr.Index.(*ast.LiteralString)
-				if !indexIsString {
-					continue
-				}
-				obj := findObjectFieldInObject(containerObj, indexValue.Value)
-				if obj != nil {
-					stack.Push(obj.Body)
-				}
-			}
-		case *ast.Var:
-			varReference, err := findVarReference(curr, vm)
-			if err != nil {
-				log.WithError(err).Errorf("Error finding var reference, ignoring this node")
-				continue
-			}
-			stack.Push(varReference)
-		}
-	}
-	return objects
-}
-
 // findVarReference finds the object that the variable is referencing
 // To do so, we get the stack where the var is used and search that stack for the var's definition
 func findVarReference(varNode *ast.Var, vm *jsonnet.VM) (ast.Node, error) {
@@ -274,7 +240,7 @@ func findVarReference(varNode *ast.Var, vm *jsonnet.VM) (ast.Node, error) {
 
 func findLhsDesugaredObject(stack *nodestack.NodeStack) (*ast.DesugaredObject, error) {
 	for !stack.IsEmpty() {
-		_, curr := stack.Pop()
+		curr := stack.Pop()
 		switch curr := curr.(type) {
 		case *ast.Binary:
 			lhsNode := curr.Left
@@ -289,10 +255,10 @@ func findLhsDesugaredObject(stack *nodestack.NodeStack) (*ast.DesugaredObject, e
 			}
 		case *ast.Local:
 			for _, bind := range curr.Binds {
-				stack = stack.Push(bind.Body)
+				stack.Push(bind.Body)
 			}
 			if curr.Body != nil {
-				stack = stack.Push(curr.Body)
+				stack.Push(curr.Body)
 			}
 		}
 	}

pkg/processing/find_param.go

@@ -6,18 +6,14 @@ import (
 )
 
 func FindParameterByIdViaStack(stack *nodestack.NodeStack, id ast.Identifier) *ast.Parameter {
-	stack = stack.Clone()
-	for !stack.IsEmpty() {
-		_, curr := stack.Pop()
-		switch curr := curr.(type) {
-		case *ast.Function:
-			for _, param := range curr.Parameters {
+	for _, node := range stack.Stack {
+		if f, ok := node.(*ast.Function); ok {
+			for _, param := range f.Parameters {
 				if param.Name == id {
 					return &param
 				}
 			}
 		}
-
 	}
 	return nil
 }

pkg/processing/find_position.go

@@ -19,77 +19,77 @@ func FindNodeByPosition(node ast.Node, location ast.Location) (*nodestack.NodeSt
 	searchStack := &nodestack.NodeStack{From: stack.From}
 	var curr ast.Node
 	for !stack.IsEmpty() {
-		stack, curr = stack.Pop()
+		curr = stack.Pop()
 		// This is needed because SuperIndex only spans "key: super" and not the ".foo" after. This only occurs
 		// when super only has 1 additional index. "super.foo.bar" will not have this issue
 		if curr, isType := curr.(*ast.SuperIndex); isType {
 			curr.Loc().End.Column = curr.Loc().End.Column + len(curr.Index.(*ast.LiteralString).Value) + 1
 		}
 		inRange := position.InRange(location, *curr.Loc())
 		if inRange {
-			searchStack = searchStack.Push(curr)
+			searchStack.Push(curr)
 		} else if curr.Loc().End.IsSet() {
 			continue
 		}
 		switch curr := curr.(type) {
 		case *ast.Local:
 			for _, bind := range curr.Binds {
-				stack = stack.Push(bind.Body)
+				stack.Push(bind.Body)
 			}
 			if curr.Body != nil {
-				stack = stack.Push(curr.Body)
+				stack.Push(curr.Body)
 			}
 		case *ast.DesugaredObject:
 			for _, field := range curr.Fields {
 				body := field.Body
 				// Functions do not have a LocRange, so we use the one from the field's body
 				if funcBody, isFunc := body.(*ast.Function); isFunc {
 					funcBody.LocRange = field.LocRange
-					stack = stack.Push(funcBody)
+					stack.Push(funcBody)
 				} else {
-					stack = stack.Push(body)
+					stack.Push(body)
 				}
 			}
 			for _, local := range curr.Locals {
-				stack = stack.Push(local.Body)
+				stack.Push(local.Body)
 			}
 		case *ast.Binary:
-			stack = stack.Push(curr.Left)
-			stack = stack.Push(curr.Right)
+			stack.Push(curr.Left)
+			stack.Push(curr.Right)
 		case *ast.Array:
 			for _, element := range curr.Elements {
-				stack = stack.Push(element.Expr)
+				stack.Push(element.Expr)
 			}
 		case *ast.Apply:
 			for _, posArg := range curr.Arguments.Positional {
-				stack = stack.Push(posArg.Expr)
+				stack.Push(posArg.Expr)
 			}
 			for _, namedArg := range curr.Arguments.Named {
-				stack = stack.Push(namedArg.Arg)
+				stack.Push(namedArg.Arg)
 			}
-			stack = stack.Push(curr.Target)
+			stack.Push(curr.Target)
 		case *ast.Conditional:
-			stack = stack.Push(curr.Cond)
-			stack = stack.Push(curr.BranchTrue)
-			stack = stack.Push(curr.BranchFalse)
+			stack.Push(curr.Cond)
+			stack.Push(curr.BranchTrue)
+			stack.Push(curr.BranchFalse)
 		case *ast.Error:
-			stack = stack.Push(curr.Expr)
+			stack.Push(curr.Expr)
 		case *ast.Function:
 			for _, param := range curr.Parameters {
 				if param.DefaultArg != nil {
-					stack = stack.Push(param.DefaultArg)
+					stack.Push(param.DefaultArg)
 				}
 			}
-			stack = stack.Push(curr.Body)
+			stack.Push(curr.Body)
 		case *ast.Index:
-			stack = stack.Push(curr.Target)
-			stack = stack.Push(curr.Index)
+			stack.Push(curr.Target)
+			stack.Push(curr.Index)
 		case *ast.InSuper:
-			stack = stack.Push(curr.Index)
+			stack.Push(curr.Index)
 		case *ast.SuperIndex:
-			stack = stack.Push(curr.Index)
+			stack.Push(curr.Index)
 		case *ast.Unary:
-			stack = stack.Push(curr.Expr)
+			stack.Push(curr.Expr)
 		}
 	}
 	return searchStack.ReorderDesugaredObjects(), nil