Merge pull request #199 from vmarkovtsev/uastv2

Port to bblfsh sdk v3

.appveyor.yml

@@ -8,7 +8,7 @@ environment:
   GOPATH: c:\gopath
 
 install:
-  - curl -SLko protoc.zip https://github.com/google/protobuf/releases/download/v3.5.1/protoc-3.5.1-win32.zip
+  - curl -SLko protoc.zip https://github.com/google/protobuf/releases/download/v3.6.0/protoc-3.6.0-win32.zip
   - 7z e protoc.zip
   - move protoc.exe C:\msys64\mingw64\bin
 

.travis.yml

@@ -13,11 +13,6 @@ go:
 services:
   - docker
 
-addons:
-  apt:
-    packages:
-    - libxml2-dev
-
 cache:
   directories:
     - $HOME/.cache/pip

Gopkg.lock

@@ -110,6 +110,14 @@
 
 [[projects]]
   branch = "master"
+  digest = "1:1a1206efd03a54d336dce7bb8719e74f2f8932f661cb9f57d5813a1d99c083d8"
+  name = "github.com/grpc-ecosystem/grpc-opentracing"
+  packages = ["go/otgrpc"]
+  pruneopts = "UT"
+  revision = "8e809c8a86450a29b90dcc9efbf062d0fe6d9746"
+
+[[projects]]
+  branch = "master"
   digest = "1:82ae4a01ab513802a68d4094e0eca6c5aaa9445633560e2c0491263e37361620"
   name = "github.com/huandu/xstrings"
   packages = ["."]
@@ -197,6 +205,18 @@
   version = "6.3.107"
 
 [[projects]]
+  digest = "1:450b7623b185031f3a456801155c8320209f75d0d4c4e633c6b1e59d44d6e392"
+  name = "github.com/opentracing/opentracing-go"
+  packages = [
+    ".",
+    "ext",
+    "log",
+  ]
+  pruneopts = "UT"
+  revision = "1949ddbfd147afd4d964a9f00b24eb291e0e7c38"
+  version = "v1.0.2"
+
+[[projects]]
   branch = "master"
   digest = "1:cab1f682826733526a4b49a45cb64fab31930f2cc2684fa14c64b090b1c25b13"
   name = "github.com/pelletier/go-buffruneio"
@@ -418,6 +438,17 @@
   version = "v2.8.8"
 
 [[projects]]
+  digest = "1:fafaec170283195ece36f81308d93913f10d5b650302d686ba28789cf70b6d01"
+  name = "gopkg.in/bblfsh/client-go.v3"
+  packages = [
+    ".",
+    "tools",
+  ]
+  pruneopts = "UT"
+  revision = "f1bb4239c6001f5f0b1f370b7d4a64866464c07b"
+  version = "v3.2.0"
+
+[[projects]]
   digest = "1:861bb28098fdd02360af9d9a1c246cce4e255c94d6a75661e74dd60144d72105"
   name = "gopkg.in/bblfsh/sdk.v1"
   packages = [
@@ -430,10 +461,11 @@
   version = "v1.16.1"
 
 [[projects]]
-  digest = "1:6346689fd052f5857391860ca459e3e8a53db2fe5da5165d8be8ec17735bd4c0"
+  digest = "1:54b8e8f6199fa36245d5bb3f96cb72b4e2d463258256c3a0e34857df45d03890"
   name = "gopkg.in/bblfsh/sdk.v2"
   packages = [
     "driver",
+    "driver/errors",
     "driver/manifest",
     "protocol",
     "protocol/v1",
@@ -447,8 +479,8 @@
     "uast/transformer",
   ]
   pruneopts = "UT"
-  revision = "73dccf5b138866d38733c4b371fcd597e4240516"
-  version = "v2.3.0"
+  revision = "e39a01efecd8a8670471352b3276c17d5e0dd7a8"
+  version = "v2.14.1"
 
 [[projects]]
   branch = "master"
@@ -615,8 +647,14 @@
     "golang.org/x/crypto/ssh/terminal",
     "gopkg.in/bblfsh/client-go.v2",
     "gopkg.in/bblfsh/client-go.v2/tools",
-    "gopkg.in/bblfsh/sdk.v1/protocol",
+    "gopkg.in/bblfsh/client-go.v3",
+    "gopkg.in/bblfsh/client-go.v3/tools",
     "gopkg.in/bblfsh/sdk.v1/uast",
+    "gopkg.in/bblfsh/sdk.v2/uast",
+    "gopkg.in/bblfsh/sdk.v2/uast/nodes",
+    "gopkg.in/bblfsh/sdk.v2/uast/nodes/nodesproto",
+    "gopkg.in/bblfsh/sdk.v2/uast/query",
+    "gopkg.in/bblfsh/sdk.v2/uast/role",
     "gopkg.in/cheggaaa/pb.v1",
     "gopkg.in/src-d/enry.v1",
     "gopkg.in/src-d/go-billy-siva.v4",

Gopkg.toml

@@ -55,12 +55,12 @@
   name = "golang.org/x/crypto"
 
 [[constraint]]
-  name = "gopkg.in/bblfsh/client-go.v2"
-  version = "2.8.8"
+  name = "gopkg.in/bblfsh/client-go.v3"
+  version = "3.2.0"
 
 [[constraint]]
-  name = "gopkg.in/bblfsh/sdk.v1"
-  version = "1.16.1"
+  name = "gopkg.in/bblfsh/sdk.v2"
+  version = "2.14.1"
 
 [[constraint]]
   branch = "master"

Makefile

@@ -6,7 +6,6 @@ EXE = .exe
 endif
 PKG = $(shell go env GOOS)_$(shell go env GOARCH)
 TAGS ?=
-BBLFSH_DEP =
 
 all: ${GOPATH}/bin/hercules${EXE}
 
@@ -34,12 +33,5 @@ cmd/hercules/plugin_template_source.go: cmd/hercules/plugin.template
 vendor:
 	dep ensure -v
 
-ifeq ($(OS),Windows_NT)
-BBLFSH_DEP = vendor/gopkg.in/bblfsh/client-go.v2/tools/include
-
-vendor/gopkg.in/bblfsh/client-go.v2/tools/include:
-	cd vendor/gopkg.in/bblfsh/client-go.v2 && make cgo-dependencies
-endif
-
-${GOPATH}/bin/hercules${EXE}: vendor *.go */*.go */*/*.go */*/*/*.go internal/pb/pb.pb.go internal/pb/pb_pb2.py cmd/hercules/plugin_template_source.go ${BBLFSH_DEP}
+${GOPATH}/bin/hercules${EXE}: vendor *.go */*.go */*/*.go */*/*/*.go internal/pb/pb.pb.go internal/pb/pb_pb2.py cmd/hercules/plugin_template_source.go
 	go get -tags "$(TAGS)" -ldflags "-X gopkg.in/src-d/hercules.v7.BinaryGitHash=$(shell git rev-parse HEAD)" gopkg.in/src-d/hercules.v7/cmd/hercules

internal/plumbing/uast/changes_xpather.go

@@ -6,8 +6,9 @@ import (
 	"log"
 
 	"github.com/minio/highwayhash"
-	"gopkg.in/bblfsh/client-go.v2/tools"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/client-go.v3/tools"
+	bblfsh "gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 )
 
@@ -22,8 +23,8 @@ var hashKey = []byte{
 }
 
 // Extract returns the list of new or changed UAST nodes filtered by XPath.
-func (xpather ChangesXPather) Extract(changes []Change) []*uast.Node {
-	result := []*uast.Node{}
+func (xpather ChangesXPather) Extract(changes []Change) []nodes.Node {
+	var result []nodes.Node
 	for _, change := range changes {
 		if change.After == nil {
 			continue
@@ -44,21 +45,25 @@ func (xpather ChangesXPather) Extract(changes []Change) []*uast.Node {
 	return result
 }
 
-func (xpather ChangesXPather) filter(root *uast.Node, origin plumbing.Hash) []*uast.Node {
-	if root != nil {
-		nodes, err := tools.Filter(root, xpather.XPath)
-		if err != nil {
-			log.Printf("libuast filter error on object %s: %v", origin.String(), err)
-			return []*uast.Node{}
-		}
-		return nodes
+func (xpather ChangesXPather) filter(root nodes.Node, origin plumbing.Hash) []nodes.Node {
+	if root == nil {
+		return nil
+	}
+	filtered, err := tools.Filter(root, xpather.XPath)
+	if err != nil {
+		log.Printf("libuast filter error on object %s: %v", origin.String(), err)
+		return []nodes.Node{}
 	}
-	return []*uast.Node{}
+	var result []nodes.Node
+	for filtered.Next() {
+		result = append(result, filtered.Node().(nodes.Node))
+	}
+	return result
 }
 
-func (xpather ChangesXPather) hash(nodes []*uast.Node) map[uint64]*uast.Node {
-	result := map[uint64]*uast.Node{}
-	for _, node := range nodes {
+func (xpather ChangesXPather) hash(nodesToHash []nodes.Node) map[uint64]nodes.Node {
+	result := map[uint64]nodes.Node{}
+	for _, node := range nodesToHash {
 		buffer := &bytes.Buffer{}
 		stringifyUASTNode(node, buffer)
 		result[highwayhash.Sum64(buffer.Bytes(), hashKey)] = node
@@ -66,9 +71,8 @@ func (xpather ChangesXPather) hash(nodes []*uast.Node) map[uint64]*uast.Node {
 	return result
 }
 
-func stringifyUASTNode(node *uast.Node, writer io.Writer) {
-	writer.Write([]byte(node.Token + "|" + node.InternalType + ">"))
-	for _, child := range node.Children {
-		stringifyUASTNode(child, writer)
+func stringifyUASTNode(node nodes.Node, writer io.Writer) {
+	for element := range tools.Iterate(tools.NewIterator(node, tools.PositionOrder)) {
+		writer.Write([]byte(bblfsh.TokenOf(element.(nodes.Object)) + "|" + bblfsh.TypeOf(element) + ">"))
 	}
 }

internal/plumbing/uast/changes_xpather_test.go

@@ -7,7 +7,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/bblfsh/client-go.v2"
+	"gopkg.in/bblfsh/client-go.v3"
 	uast_test "gopkg.in/src-d/hercules.v7/internal/plumbing/uast/test"
 	"gopkg.in/src-d/hercules.v7/internal/test"
 )
@@ -27,7 +27,7 @@ func TestChangesXPatherExtractChanged(t *testing.T) {
 		{Before: nil, After: root2, Change: gitChange},
 		{Before: root1, After: nil, Change: gitChange},
 	}
-	xpather := ChangesXPather{XPath: "//*[@roleComment]"}
+	xpather := ChangesXPather{XPath: "//*[@role='Comment']"}
 	nodes := xpather.Extract(uastChanges)
 	assert.True(t, len(nodes) > 0)
 }

internal/plumbing/uast/diff_refiner.go

@@ -4,7 +4,9 @@ import (
 	"unicode/utf8"
 
 	"github.com/sergi/go-diff/diffmatchpatch"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/client-go.v3/tools"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/hercules.v7/internal/core"
 	"gopkg.in/src-d/hercules.v7/internal/plumbing"
@@ -108,15 +110,14 @@ func (ref *FileDiffRefiner) Consume(deps map[string]interface{}) (map[string]int
 			result[fileName] = oldDiff
 			continue
 		}
-		line2node := make([][]*uast.Node, oldDiff.NewLinesOfCode)
-		VisitEachNode(uastChange.After, func(node *uast.Node) {
-			if node.StartPosition != nil && node.EndPosition != nil {
-				for l := node.StartPosition.Line; l <= node.EndPosition.Line; l++ {
-					nodes := line2node[l-1] // line starts with 1
-					if nodes == nil {
-						nodes = []*uast.Node{}
+		line2node := make([][]nodes.Node, oldDiff.NewLinesOfCode)
+		VisitEachNode(uastChange.After, func(node nodes.Node) {
+			if obj, ok := node.(nodes.Object); ok {
+				pos := uast.PositionsOf(obj)
+				if pos.Start() != nil && pos.End() != nil {
+					for l := pos.Start().Line; l <= pos.End().Line; l++ {
+						line2node[l-1] = append(line2node[l-1], node) // line starts with 1
 					}
-					line2node[l-1] = append(nodes, node)
 				}
 			}
 		})
@@ -172,27 +173,22 @@ func (ref *FileDiffRefiner) Fork(n int) []core.PipelineItem {
 
 // VisitEachNode is a handy routine to execute a callback on every node in the subtree,
 // including the root itself. Depth first tree traversal.
-func VisitEachNode(root *uast.Node, payload func(*uast.Node)) {
-	queue := []*uast.Node{}
-	queue = append(queue, root)
-	for len(queue) > 0 {
-		node := queue[len(queue)-1]
-		queue = queue[:len(queue)-1]
-		payload(node)
-		for _, child := range node.Children {
-			queue = append(queue, child)
+func VisitEachNode(root nodes.Node, payload func(nodes.Node)) {
+	for child := range tools.Iterate(tools.NewIterator(root, tools.PreOrder)) {
+		if _, ok := child.(nodes.Object); ok {
+			payload(child)
 		}
 	}
 }
 
-func countNodesInInterval(occupiedMap [][]*uast.Node, start, end int) int {
-	nodes := map[*uast.Node]bool{}
+func countNodesInInterval(occupiedMap [][]nodes.Node, start, end int) int {
+	inodes := map[nodes.Comparable]bool{}
 	for i := start; i < end; i++ {
 		for _, node := range occupiedMap[i] {
-			nodes[node] = true
+			inodes[nodes.UniqueKey(node)] = true
 		}
 	}
-	return len(nodes)
+	return len(inodes)
 }
 
 func init() {

internal/plumbing/uast/diff_refiner_test.go

@@ -2,14 +2,15 @@ package uast
 
 import (
 	"io/ioutil"
+	"os"
 	"path"
 	"testing"
 	"unicode/utf8"
 
-	"github.com/gogo/protobuf/proto"
 	"github.com/sergi/go-diff/diffmatchpatch"
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes/nodesproto"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
 	"gopkg.in/src-d/hercules.v7/internal/core"
 	"gopkg.in/src-d/hercules.v7/internal/plumbing"
@@ -50,6 +51,19 @@ func TestFileDiffRefinerRegistration(t *testing.T) {
 	assert.True(t, matched)
 }
 
+func loadUast(t *testing.T, name string) nodes.Node {
+	filename := path.Join("..", "..", "test_data", name)
+	reader, err := os.Open(filename)
+	if err != nil {
+		assert.Failf(t, "cannot load %s: %v", filename, err)
+	}
+	node, err := nodesproto.ReadTree(reader)
+	if err != nil {
+		assert.Failf(t, "cannot load %s: %v", filename, err)
+	}
+	return node
+}
+
 func TestFileDiffRefinerConsume(t *testing.T) {
 	bytes1, err := ioutil.ReadFile(path.Join("..", "..", "test_data", "1.java"))
 	assert.Nil(t, err)
@@ -67,19 +81,12 @@ func TestFileDiffRefinerConsume(t *testing.T) {
 	}
 	state[plumbing.DependencyFileDiff] = fileDiffs
 	uastChanges := make([]Change, 1)
-	loadUast := func(name string) *uast.Node {
-		bytes, err := ioutil.ReadFile(path.Join("..", "..", "test_data", name))
-		assert.Nil(t, err)
-		node := uast.Node{}
-		proto.Unmarshal(bytes, &node)
-		return &node
-	}
 	state[DependencyUastChanges] = uastChanges
 	uastChanges[0] = Change{
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: fileName},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: loadUast("uast1.pb"), After: loadUast("uast2.pb"),
+		Before: loadUast(t, "uast1.pb"), After: loadUast(t, "uast2.pb"),
 	}
 	fd := fixtureFileDiffRefiner()
 	iresult, err := fd.Consume(state)
@@ -116,19 +123,12 @@ func TestFileDiffRefinerConsumeNoUast(t *testing.T) {
 	}
 	state[plumbing.DependencyFileDiff] = fileDiffs
 	uastChanges := make([]Change, 1)
-	loadUast := func(name string) *uast.Node {
-		bytes, err := ioutil.ReadFile(path.Join("..", "..", "test_data", name))
-		assert.Nil(t, err)
-		node := uast.Node{}
-		proto.Unmarshal(bytes, &node)
-		return &node
-	}
 	state[DependencyUastChanges] = uastChanges
 	uastChanges[0] = Change{
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: fileName},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: loadUast("uast1.pb"), After: nil,
+		Before: loadUast(t, "uast1.pb"), After: nil,
 	}
 	fd := fixtureFileDiffRefiner()
 	iresult, err := fd.Consume(state)
@@ -145,7 +145,7 @@ func TestFileDiffRefinerConsumeNoUast(t *testing.T) {
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: fileName},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: loadUast("uast1.pb"), After: loadUast("uast2.pb"),
+		Before: loadUast(t, "uast1.pb"), After: loadUast(t, "uast2.pb"),
 	}
 	iresult, err = fd.Consume(state)
 	assert.Nil(t, err)

internal/plumbing/uast/test/utils.go

@@ -1,17 +1,16 @@
 package test
 
 import (
-	"fmt"
 	"io/ioutil"
 
-	"gopkg.in/bblfsh/client-go.v2"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/client-go.v3"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	core_test "gopkg.in/src-d/hercules.v7/internal/test"
 )
 
 // ParseBlobFromTestRepo extracts the UAST from the file by it's hash and name.
-func ParseBlobFromTestRepo(hash, name string, client *bblfsh.Client) *uast.Node {
+func ParseBlobFromTestRepo(hash, name string, client *bblfsh.Client) nodes.Node {
 	blob, err := core_test.Repository.BlobObject(plumbing.NewHash(hash))
 	if err != nil {
 		panic(err)
@@ -25,15 +24,10 @@ func ParseBlobFromTestRepo(hash, name string, client *bblfsh.Client) *uast.Node
 	if err != nil {
 		panic(err)
 	}
-	request := client.NewParseRequest()
-	request.Content(string(data))
-	request.Filename(name)
-	response, err := request.Do()
+	request := client.NewParseRequest().Content(string(data)).Filename(name).Mode(bblfsh.Annotated)
+	response, _, err := request.UAST()
 	if err != nil {
 		panic(err)
 	}
-	if response.UAST == nil {
-		panic(fmt.Sprintf("empty response for %s %s", name, hash))
-	}
-	return response.UAST
+	return response
 }

internal/plumbing/uast/uast.go

@@ -16,9 +16,9 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/jeffail/tunny"
-	"gopkg.in/bblfsh/client-go.v2"
-	"gopkg.in/bblfsh/sdk.v1/protocol"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/client-go.v3"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes/nodesproto"
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
@@ -63,7 +63,7 @@ const (
 
 type uastTask struct {
 	Lock   *sync.RWMutex
-	Dest   map[plumbing.Hash]*uast.Node
+	Dest   map[plumbing.Hash]nodes.Node
 	Name   string
 	Hash   plumbing.Hash
 	Data   []byte
@@ -207,7 +207,7 @@ func (exr *Extractor) Initialize(repository *git.Repository) error {
 func (exr *Extractor) Consume(deps map[string]interface{}) (map[string]interface{}, error) {
 	cache := deps[items.DependencyBlobCache].(map[plumbing.Hash]*items.CachedBlob)
 	treeDiffs := deps[items.DependencyTreeChanges].(object.Changes)
-	uasts := map[plumbing.Hash]*uast.Node{}
+	uasts := map[plumbing.Hash]nodes.Node{}
 	lock := sync.RWMutex{}
 	errs := make([]error, 0)
 	wg := sync.WaitGroup{}
@@ -261,28 +261,21 @@ func (exr *Extractor) Fork(n int) []core.PipelineItem {
 }
 
 func (exr *Extractor) extractUAST(
-	client *bblfsh.Client, name string, data []byte) (*uast.Node, error) {
-	request := client.NewParseRequest()
-	request.Content(string(data))
-	request.Filename(name)
+	client *bblfsh.Client, name string, data []byte) (nodes.Node, error) {
 	ctx, cancel := exr.Context()
 	if cancel != nil {
 		defer cancel()
 	}
-	response, err := request.DoWithContext(ctx)
+	request := client.NewParseRequest().
+		Content(string(data)).Filename(name).Mode(bblfsh.Annotated).Context(ctx)
+	response, _, err := request.UAST()
 	if err != nil {
 		if strings.Contains("missing driver", err.Error()) {
 			return nil, nil
 		}
 		return nil, err
 	}
-	if response.Status != protocol.Ok {
-		return nil, errors.New(strings.Join(response.Errors, "\n"))
-	}
-	if err != nil {
-		return nil, err
-	}
-	return response.UAST, nil
+	return response, nil
 }
 
 func (exr *Extractor) extractTask(client *bblfsh.Client, data interface{}) interface{} {
@@ -303,8 +296,8 @@ func (exr *Extractor) extractTask(client *bblfsh.Client, data interface{}) inter
 
 // Change is the type of the items in the list of changes which is provided by Changes.
 type Change struct {
-	Before *uast.Node
-	After  *uast.Node
+	Before nodes.Node
+	After  nodes.Node
 	Change *object.Change
 }
 
@@ -317,7 +310,7 @@ const (
 // in a commit. It is a PipelineItem.
 type Changes struct {
 	core.NoopMerger
-	cache map[plumbing.Hash]*uast.Node
+	cache map[plumbing.Hash]nodes.Node
 }
 
 // Name of this PipelineItem. Uniquely identifies the type, used for mapping keys, etc.
@@ -360,7 +353,7 @@ func (uc *Changes) Configure(facts map[string]interface{}) error {
 // Initialize resets the temporary caches and prepares this PipelineItem for a series of Consume()
 // calls. The repository which is going to be analysed is supplied as an argument.
 func (uc *Changes) Initialize(repository *git.Repository) error {
-	uc.cache = map[plumbing.Hash]*uast.Node{}
+	uc.cache = map[plumbing.Hash]nodes.Node{}
 	return nil
 }
 
@@ -370,7 +363,7 @@ func (uc *Changes) Initialize(repository *git.Repository) error {
 // This function returns the mapping with analysis results. The keys must be the same as
 // in Provides(). If there was an error, nil is returned.
 func (uc *Changes) Consume(deps map[string]interface{}) (map[string]interface{}, error) {
-	uasts := deps[DependencyUasts].(map[plumbing.Hash]*uast.Node)
+	uasts := deps[DependencyUasts].(map[plumbing.Hash]nodes.Node)
 	treeDiffs := deps[items.DependencyTreeChanges].(object.Changes)
 	commit := make([]Change, 0, len(treeDiffs))
 	for _, change := range treeDiffs {
@@ -405,7 +398,7 @@ func (uc *Changes) Fork(n int) []core.PipelineItem {
 	ucs := make([]core.PipelineItem, n)
 	for i := 0; i < n; i++ {
 		clone := &Changes{
-			cache: map[plumbing.Hash]*uast.Node{},
+			cache: map[plumbing.Hash]nodes.Node{},
 		}
 		for key, val := range uc.cache {
 			clone.cache[key] = val
@@ -535,31 +528,46 @@ func (saver *ChangesSaver) Serialize(result interface{}, binary bool, writer io.
 }
 
 func (saver *ChangesSaver) dumpFiles(result [][]Change) []*pb.UASTChange {
-	fileNames := []*pb.UASTChange{}
+	var fileNames []*pb.UASTChange
+	dumpUast := func(uast nodes.Node, path string) {
+		f, err := os.Create(path)
+		if err != nil {
+			panic(err)
+		}
+		defer f.Close()
+		err = nodesproto.WriteTo(f, uast)
+		if err != nil {
+			panic(err)
+		}
+	}
 	for i, changes := range result {
 		for j, change := range changes {
 			if change.Before == nil || change.After == nil {
 				continue
 			}
 			record := &pb.UASTChange{FileName: change.Change.To.Name}
-			bs, _ := change.Before.Marshal()
 			record.UastBefore = path.Join(saver.OutputPath, fmt.Sprintf(
 				"%d_%d_before_%s.pb", i, j, change.Change.From.TreeEntry.Hash.String()))
-			ioutil.WriteFile(record.UastBefore, bs, 0666)
+			dumpUast(change.Before, record.UastBefore)
 			blob, _ := saver.repository.BlobObject(change.Change.From.TreeEntry.Hash)
 			s, _ := (&object.File{Blob: *blob}).Contents()
 			record.SrcBefore = path.Join(saver.OutputPath, fmt.Sprintf(
 				"%d_%d_before_%s.src", i, j, change.Change.From.TreeEntry.Hash.String()))
-			ioutil.WriteFile(record.SrcBefore, []byte(s), 0666)
-			bs, _ = change.After.Marshal()
+			err := ioutil.WriteFile(record.SrcBefore, []byte(s), 0666)
+			if err != nil {
+				panic(err)
+			}
 			record.UastAfter = path.Join(saver.OutputPath, fmt.Sprintf(
 				"%d_%d_after_%s.pb", i, j, change.Change.To.TreeEntry.Hash.String()))
-			ioutil.WriteFile(record.UastAfter, bs, 0666)
+			dumpUast(change.After, record.UastAfter)
 			blob, _ = saver.repository.BlobObject(change.Change.To.TreeEntry.Hash)
 			s, _ = (&object.File{Blob: *blob}).Contents()
 			record.SrcAfter = path.Join(saver.OutputPath, fmt.Sprintf(
 				"%d_%d_after_%s.src", i, j, change.Change.To.TreeEntry.Hash.String()))
-			ioutil.WriteFile(record.SrcAfter, []byte(s), 0666)
+			err = ioutil.WriteFile(record.SrcAfter, []byte(s), 0666)
+			if err != nil {
+				panic(err)
+			}
 			fileNames = append(fileNames, record)
 		}
 	}

internal/plumbing/uast/uast_test.go

@@ -12,7 +12,8 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
 	"gopkg.in/src-d/hercules.v7/internal/core"
@@ -172,9 +173,9 @@ func TestUASTExtractorConsume(t *testing.T) {
 
 	res, err = exr.Consume(deps)
 	assert.Nil(t, err)
-	uasts := res[DependencyUasts].(map[plumbing.Hash]*uast.Node)
+	uasts := res[DependencyUasts].(map[plumbing.Hash]nodes.Node)
 	assert.Equal(t, len(uasts), 1)
-	assert.Equal(t, len(uasts[hash].Children), 24)
+	assert.Equal(t, len(uasts[hash].(nodes.Object)["body"].(nodes.Array)), 24)
 }
 
 func TestUASTExtractorFork(t *testing.T) {
@@ -221,24 +222,27 @@ func TestUASTChangesRegistration(t *testing.T) {
 	assert.True(t, matched)
 }
 
+func newNodeWithType(name string) nodes.Node {
+	return nodes.Object{
+		uast.KeyType:  nodes.String(name),
+		uast.KeyToken: nodes.String("my_token"),
+	}
+}
+
 func TestUASTChangesConsume(t *testing.T) {
-	var uastsArray []*uast.Node
-	uasts := map[plumbing.Hash]*uast.Node{}
+	var uastsArray []nodes.Node
+	uasts := map[plumbing.Hash]nodes.Node{}
 	hash := plumbing.NewHash("291286b4ac41952cbd1389fda66420ec03c1a9fe")
-	uasts[hash] = &uast.Node{}
-	uasts[hash].InternalType = "uno"
+	uasts[hash] = newNodeWithType("uno")
 	uastsArray = append(uastsArray, uasts[hash])
 	hash = plumbing.NewHash("c29112dbd697ad9b401333b80c18a63951bc18d9")
-	uasts[hash] = &uast.Node{}
-	uasts[hash].InternalType = "dos"
+	uasts[hash] = newNodeWithType("dos")
 	uastsArray = append(uastsArray, uasts[hash])
 	hash = plumbing.NewHash("baa64828831d174f40140e4b3cfa77d1e917a2c1")
-	uasts[hash] = &uast.Node{}
-	uasts[hash].InternalType = "tres"
+	uasts[hash] = newNodeWithType("tres")
 	uastsArray = append(uastsArray, uasts[hash])
 	hash = plumbing.NewHash("dc248ba2b22048cc730c571a748e8ffcf7085ab9")
-	uasts[hash] = &uast.Node{}
-	uasts[hash].InternalType = "quatro"
+	uasts[hash] = newNodeWithType("quatro")
 	uastsArray = append(uastsArray, uasts[hash])
 	changes := make(object.Changes, 3)
 	treeFrom, _ := test.Repository.TreeObject(plumbing.NewHash(
@@ -374,7 +378,7 @@ func TestUASTChangesSaverPayload(t *testing.T) {
 		"a1eb2ea76eb7f9bfbde9b243861474421000eb96"))
 	treeTo, _ := test.Repository.TreeObject(plumbing.NewHash(
 		"994eac1cd07235bb9815e547a75c84265dea00f5"))
-	changes[0] = Change{Before: &uast.Node{}, After: &uast.Node{},
+	changes[0] = Change{Before: nodes.Object{}, After: nodes.Object{},
 		Change: &object.Change{From: object.ChangeEntry{
 			Name: "analyser.go",
 			Tree: treeFrom,
@@ -448,7 +452,7 @@ func TestUASTChangesSaverConsumeMerge(t *testing.T) {
 		"a1eb2ea76eb7f9bfbde9b243861474421000eb96"))
 	treeTo, _ := test.Repository.TreeObject(plumbing.NewHash(
 		"994eac1cd07235bb9815e547a75c84265dea00f5"))
-	changes[0] = Change{Before: &uast.Node{}, After: &uast.Node{},
+	changes[0] = Change{Before: nodes.Object{}, After: nodes.Object{},
 		Change: &object.Change{From: object.ChangeEntry{
 			Name: "analyser.go",
 			Tree: treeFrom,

leaves/comment_sentiment.go

@@ -12,7 +12,8 @@ import (
 	"strings"
 
 	"github.com/gogo/protobuf/proto"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
 	progress "gopkg.in/cheggaaa/pb.v1"
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/go-git.v4/plumbing"
@@ -149,7 +150,7 @@ func (sent *CommentSentimentAnalysis) validate() {
 // calls. The repository which is going to be analysed is supplied as an argument.
 func (sent *CommentSentimentAnalysis) Initialize(repository *git.Repository) error {
 	sent.commentsByDay = map[int][]string{}
-	sent.xpather = &uast_items.ChangesXPather{XPath: "//*[@roleComment]"}
+	sent.xpather = &uast_items.ChangesXPather{XPath: "//*[@role='Comment']"}
 	sent.validate()
 	sent.OneShotMergeProcessor.Initialize()
 	return nil
@@ -298,20 +299,16 @@ func (sent *CommentSentimentAnalysis) serializeBinary(
 	return nil
 }
 
-func (sent *CommentSentimentAnalysis) mergeComments(nodes []*uast.Node) []string {
+func (sent *CommentSentimentAnalysis) mergeComments(extracted []nodes.Node) []string {
 	var mergedComments []string
-	lines := map[int][]*uast.Node{}
-	for _, node := range nodes {
-		if node.StartPosition == nil {
+	lines := map[int][]nodes.Node{}
+	for _, node := range extracted {
+		pos := uast.PositionsOf(node.(nodes.Object))
+		if pos.Start() == nil {
 			continue
 		}
-		lineno := int(node.StartPosition.Line)
-		subnodes := lines[lineno]
-		if subnodes == nil {
-			subnodes = []*uast.Node{}
-		}
-		subnodes = append(subnodes, node)
-		lines[lineno] = subnodes
+		lineno := int(pos.Start().Line)
+		lines[lineno] = append(lines[lineno], node)
 	}
 	lineNums := make([]int, 0, len(lines))
 	for line := range lines {
@@ -323,10 +320,15 @@ func (sent *CommentSentimentAnalysis) mergeComments(nodes []*uast.Node) []string
 		lineNodes := lines[line]
 		maxEnd := line
 		for _, node := range lineNodes {
-			if node.EndPosition != nil && maxEnd < int(node.EndPosition.Line) {
-				maxEnd = int(node.EndPosition.Line)
+			pos := uast.PositionsOf(node.(nodes.Object))
+			if pos.End() != nil && maxEnd < int(pos.End().Line) {
+				maxEnd = int(pos.End().Line)
+			}
+			token := strings.TrimSpace(uast.TokenOf(node.(nodes.Object)))
+			// FIXME(vmarkovtsev): remove this hack when https://github.com/bblfsh/go-driver/issues/39 is fixed
+			if len(token) > 0 && token[0] == '#' {
+				token = token[1:]
 			}
-			token := strings.TrimSpace(node.Token)
 			if token != "" {
 				buffer = append(buffer, token)
 			}

leaves/comment_sentiment_test.go

@@ -10,8 +10,11 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/bblfsh/client-go.v2"
-	"gopkg.in/bblfsh/client-go.v2/tools"
+	"gopkg.in/bblfsh/client-go.v3"
+	"gopkg.in/bblfsh/client-go.v3/tools"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
+	"gopkg.in/bblfsh/sdk.v2/uast/query"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/hercules.v7/internal/core"
 	"gopkg.in/src-d/hercules.v7/internal/pb"
@@ -153,13 +156,14 @@ func TestCommentSentimentConsume(t *testing.T) {
 	hash2 := "2a7392320b332494a08d5113aabe6d056fef7e9d"
 	root1 := uast_test.ParseBlobFromTestRepo(hash1, "labours.py", client)
 	root2 := uast_test.ParseBlobFromTestRepo(hash2, "labours.py", client)
-	comments, _ := tools.Filter(root2, "//*[@roleComment]")
-	for _, c := range comments {
-		t := strings.TrimSpace(c.Token)
+	comments, _ := tools.Filter(root2, "//*[@role='Comment']")
+	for _, c := range query.AllNodes(comments) {
+		obj := c.(nodes.Object)
+		t := strings.TrimSpace(uast.TokenOf(obj))
 		if t == "we need to adjust the peak, it may not be less than the decayed value" {
-			c.Token = "license copyright boring"
+			obj[uast.KeyToken] = nodes.String("license copyright boring")
 		} else if t == "Tensorflow 1.5 parses sys.argv unconditionally *applause*" {
-			c.StartPosition = nil
+			obj[uast.KeyPos].(nodes.Object)[uast.KeyStart] = nil
 		}
 	}
 	gitChange := test.FakeChangeForName("labours.py", hash1, hash2)
@@ -169,11 +173,13 @@ func TestCommentSentimentConsume(t *testing.T) {
 			{Before: root1, After: root2, Change: gitChange},
 		},
 	}
+	deps[core.DependencyCommit], _ = test.Repository.CommitObject(plumbing.NewHash(
+		"d48d50d0b72be95bb8054e3839b8dd79fa970b1c"))
 	result, err := sent.Consume(deps)
 	assert.Nil(t, err)
 	assert.Nil(t, result)
 	assert.Len(t, sent.commentsByDay, 1)
-	assert.Len(t, sent.commentsByDay[0], 4)
+	assert.Len(t, sent.commentsByDay[0], 6)
 }
 
 var (

leaves/shotness.go

@@ -9,8 +9,11 @@ import (
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/sergi/go-diff/diffmatchpatch"
-	"gopkg.in/bblfsh/client-go.v2/tools"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/client-go.v3/tools"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	uast_nodes "gopkg.in/bblfsh/sdk.v2/uast/nodes"
+	"gopkg.in/bblfsh/sdk.v2/uast/query"
+	"gopkg.in/bblfsh/sdk.v2/uast/role"
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
 	"gopkg.in/src-d/hercules.v7/internal/core"
@@ -42,10 +45,10 @@ const (
 
 	// DefaultShotnessXpathStruct is the default UAST XPath to choose the analysed nodes.
 	// It extracts functions.
-	DefaultShotnessXpathStruct = "//*[@roleFunction and @roleDeclaration]"
+	DefaultShotnessXpathStruct = "//*[@role='Function' and @role='Declaration']"
 	// DefaultShotnessXpathName is the default UAST XPath to choose the names of the analysed nodes.
 	// It looks at the current tree level and at the immediate children.
-	DefaultShotnessXpathName = "/*[@roleFunction and @roleIdentifier and @roleName] | /*/*[@roleFunction and @roleIdentifier and @roleName]"
+	DefaultShotnessXpathName = "/*[1]/*/*[@role='Function' and @role='Identifier' and @role='Name'] | /*[1]"
 )
 
 type nodeShotness struct {
@@ -58,7 +61,7 @@ type nodeShotness struct {
 // These attributes are supposed to uniquely identify each node.
 type NodeSummary struct {
 	InternalRole string
-	Roles        []uast.Role
+	Roles        []role.Role
 	Name         string
 	File         string
 }
@@ -167,10 +170,10 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 	diffs := deps[items.DependencyFileDiff].(map[string]items.FileDiffData)
 	allNodes := map[string]bool{}
 
-	addNode := func(name string, node *uast.Node, fileName string) {
+	addNode := func(name string, node uast_nodes.Node, fileName string) {
 		nodeSummary := NodeSummary{
-			InternalRole: node.InternalType,
-			Roles:        node.Roles,
+			InternalRole: uast.TypeOf(node),
+			Roles:        uast.RolesOf(node.(uast_nodes.Object)),
 			Name:         name,
 			File:         fileName,
 		}
@@ -260,23 +263,25 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 			continue
 		}
 		reversedNodesAfter := reverseNodeMap(nodesAfter)
-		genLine2Node := func(nodes map[string]*uast.Node, linesNum int) [][]*uast.Node {
-			res := make([][]*uast.Node, linesNum)
+		genLine2Node := func(nodes map[string]uast_nodes.Node, linesNum int) [][]uast_nodes.Node {
+			res := make([][]uast_nodes.Node, linesNum)
 			for _, node := range nodes {
-				if node.StartPosition == nil {
+				pos := uast.PositionsOf(node.(uast_nodes.Object))
+				if pos.Start() == nil {
 					continue
 				}
-				startLine := node.StartPosition.Line
-				endLine := node.StartPosition.Line
-				if node.EndPosition != nil && node.EndPosition.Line > node.StartPosition.Line {
-					endLine = node.EndPosition.Line
+				startLine := pos.Start().Line
+				endLine := pos.Start().Line
+				if pos.End() != nil && pos.End().Line > pos.Start().Line {
+					endLine = pos.End().Line
 				} else {
-					// we need to determine node.EndPosition.Line
-					uast_items.VisitEachNode(node, func(child *uast.Node) {
-						if child.StartPosition != nil {
-							candidate := child.StartPosition.Line
-							if child.EndPosition != nil {
-								candidate = child.EndPosition.Line
+					// we need to determine pos.End().Line
+					uast_items.VisitEachNode(node, func(child uast_nodes.Node) {
+						childPos := uast.PositionsOf(child.(uast_nodes.Object))
+						if childPos.Start() != nil {
+							candidate := childPos.Start().Line
+							if childPos.End() != nil {
+								candidate = childPos.End().Line
 							}
 							if candidate > endLine {
 								endLine = candidate
@@ -287,7 +292,7 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 				for l := startLine; l <= endLine; l++ {
 					lineNodes := res[l-1]
 					if lineNodes == nil {
-						lineNodes = []*uast.Node{}
+						lineNodes = []uast_nodes.Node{}
 					}
 					lineNodes = append(lineNodes, node)
 					res[l-1] = lineNodes
@@ -309,7 +314,7 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 					nodes := line2nodeBefore[l]
 					for _, node := range nodes {
 						// toName because we handled a possible rename before
-						addNode(reversedNodesBefore[node], node, toName)
+						addNode(reversedNodesBefore[uast_nodes.UniqueKey(node)], node, toName)
 					}
 				}
 				lineNumBefore += size
@@ -317,7 +322,7 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 				for l := lineNumAfter; l < lineNumAfter+size; l++ {
 					nodes := line2nodeAfter[l]
 					for _, node := range nodes {
-						addNode(reversedNodesAfter[node], node, toName)
+						addNode(reversedNodesAfter[uast_nodes.UniqueKey(node)], node, toName)
 					}
 				}
 				lineNumAfter += size
@@ -443,49 +448,54 @@ func (shotness *ShotnessAnalysis) serializeBinary(result *ShotnessResult, writer
 	return err
 }
 
-func (shotness *ShotnessAnalysis) extractNodes(root *uast.Node) (map[string]*uast.Node, error) {
-	structs, err := tools.Filter(root, shotness.XpathStruct)
+func (shotness *ShotnessAnalysis) extractNodes(root uast_nodes.Node) (map[string]uast_nodes.Node, error) {
+	it, err := tools.Filter(root, shotness.XpathStruct)
 	if err != nil {
 		return nil, err
 	}
+	structs := query.AllNodes(it)
 	// some structs may be inside other structs; we pick the outermost
 	// otherwise due to UAST quirks there may be false positives
-	internal := map[*uast.Node]bool{}
-	for _, mainNode := range structs {
-		if internal[mainNode] {
+	internal := map[uast_nodes.Comparable]bool{}
+	for _, ext := range structs {
+		mainNode := ext.(uast_nodes.Node)
+		if internal[uast_nodes.UniqueKey(mainNode)] {
 			continue
 		}
 		subs, err := tools.Filter(mainNode, shotness.XpathStruct)
 		if err != nil {
 			return nil, err
 		}
-		for _, sub := range subs {
-			if sub != mainNode {
-				internal[sub] = true
+		for subs.Next() {
+			sub := subs.Node().(uast_nodes.Node)
+			if uast_nodes.UniqueKey(sub) != uast_nodes.UniqueKey(mainNode) {
+				internal[uast_nodes.UniqueKey(sub)] = true
 			}
 		}
 	}
-	res := map[string]*uast.Node{}
-	for _, node := range structs {
-		if internal[node] {
+	res := map[string]uast_nodes.Node{}
+	for _, ext := range structs {
+		node := ext.(uast_nodes.Node)
+		if internal[uast_nodes.UniqueKey(node)] {
 			continue
 		}
-		nodeNames, err := tools.Filter(node, shotness.XpathName)
+		it, err := tools.Filter(node, shotness.XpathName)
 		if err != nil {
 			return nil, err
 		}
+		nodeNames := query.AllNodes(it)
 		if len(nodeNames) == 0 {
 			continue
 		}
-		res[nodeNames[0].Token] = node
+		res[uast.TokenOf(nodeNames[0].(uast_nodes.Object))] = node
 	}
 	return res, nil
 }
 
-func reverseNodeMap(nodes map[string]*uast.Node) map[*uast.Node]string {
-	res := map[*uast.Node]string{}
+func reverseNodeMap(nodes map[string]uast_nodes.Node) map[uast_nodes.Comparable]string {
+	res := map[uast_nodes.Comparable]string{}
 	for key, node := range nodes {
-		res[node] = key
+		res[uast_nodes.UniqueKey(node)] = key
 	}
 	return res
 }

leaves/shotness_test.go

@@ -3,13 +3,16 @@ package leaves
 import (
 	"bytes"
 	"io/ioutil"
+	"os"
 	"path"
 	"testing"
 
 	"github.com/gogo/protobuf/proto"
 	"github.com/sergi/go-diff/diffmatchpatch"
 	"github.com/stretchr/testify/assert"
-	"gopkg.in/bblfsh/sdk.v1/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes"
+	"gopkg.in/bblfsh/sdk.v2/uast/nodes/nodesproto"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
 	"gopkg.in/src-d/hercules.v7/internal/core"
 	"gopkg.in/src-d/hercules.v7/internal/pb"
@@ -67,6 +70,19 @@ func TestShotnessRegistration(t *testing.T) {
 	assert.True(t, matched)
 }
 
+func loadUast(t *testing.T, name string) nodes.Node {
+	filename := path.Join("..", "internal", "test_data", name)
+	reader, err := os.Open(filename)
+	if err != nil {
+		assert.Failf(t, "cannot load %s: %v", filename, err)
+	}
+	node, err := nodesproto.ReadTree(reader)
+	if err != nil {
+		assert.Failf(t, "cannot load %s: %v", filename, err)
+	}
+	return node
+}
+
 func bakeShotness(t *testing.T, eraseEndPosition bool) (*ShotnessAnalysis, ShotnessResult) {
 	sh := fixtureShotness()
 	bytes1, err := ioutil.ReadFile(path.Join("..", "internal", "test_data", "1.java"))
@@ -86,24 +102,25 @@ func bakeShotness(t *testing.T, eraseEndPosition bool) (*ShotnessAnalysis, Shotn
 	}
 	state[items.DependencyFileDiff] = fileDiffs
 	uastChanges := make([]uast_items.Change, 1)
-	loadUast := func(name string) *uast.Node {
-		bytes, err := ioutil.ReadFile(path.Join("..", "internal", "test_data", name))
-		assert.Nil(t, err)
-		node := uast.Node{}
-		proto.Unmarshal(bytes, &node)
+	myLoadUast := func(name string) nodes.Node {
+		node := loadUast(t, name)
 		if eraseEndPosition {
-			uast_items.VisitEachNode(&node, func(child *uast.Node) {
-				child.EndPosition = nil
+			uast_items.VisitEachNode(node, func(child nodes.Node) {
+				obj, _ := child.(nodes.Object)[uast.KeyPos].(nodes.Object)
+				if len(obj) == 0 {
+					return
+				}
+				obj[uast.KeyEnd] = nil
 			})
 		}
-		return &node
+		return node
 	}
 	state[uast_items.DependencyUastChanges] = uastChanges
 	uastChanges[0] = uast_items.Change{
 		Change: &object.Change{
 			From: object.ChangeEntry{},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: nil, After: loadUast("uast1.pb"),
+		Before: nil, After: myLoadUast("uast1.pb"),
 	}
 	iresult, err := sh.Consume(state)
 	assert.Nil(t, err)
@@ -112,7 +129,7 @@ func bakeShotness(t *testing.T, eraseEndPosition bool) (*ShotnessAnalysis, Shotn
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: fileName},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: loadUast("uast1.pb"), After: loadUast("uast2.pb"),
+		Before: myLoadUast("uast1.pb"), After: myLoadUast("uast2.pb"),
 	}
 	iresult, err = sh.Consume(state)
 	assert.Nil(t, err)
@@ -140,19 +157,12 @@ func TestShotnessConsume(t *testing.T) {
 	}
 	state[items.DependencyFileDiff] = fileDiffs
 	uastChanges := make([]uast_items.Change, 1)
-	loadUast := func(name string) *uast.Node {
-		bytes, err := ioutil.ReadFile(path.Join("..", "internal", "test_data", name))
-		assert.Nil(t, err)
-		node := uast.Node{}
-		proto.Unmarshal(bytes, &node)
-		return &node
-	}
 	state[uast_items.DependencyUastChanges] = uastChanges
 	uastChanges[0] = uast_items.Change{
 		Change: &object.Change{
 			From: object.ChangeEntry{},
 			To:   object.ChangeEntry{Name: fileName}},
-		Before: nil, After: loadUast("uast1.pb"),
+		Before: nil, After: loadUast(t, "uast1.pb"),
 	}
 	iresult, err := sh.Consume(state)
 	assert.Nil(t, err)
@@ -161,7 +171,7 @@ func TestShotnessConsume(t *testing.T) {
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: fileName},
 			To:   object.ChangeEntry{Name: newfileName}},
-		Before: loadUast("uast1.pb"), After: loadUast("uast2.pb"),
+		Before: loadUast(t, "uast1.pb"), After: loadUast(t, "uast2.pb"),
 	}
 	fileDiffs[newfileName] = fileDiffs[fileName]
 	delete(fileDiffs, fileName)
@@ -188,6 +198,9 @@ func TestShotnessConsume(t *testing.T) {
 	result := sh.Finalize().(ShotnessResult)
 	assert.Len(t, result.Nodes, 18)
 	assert.Len(t, result.Counters, 18)
+	if len(result.Nodes) != 18 || len(result.Counters) != 18 {
+		t.FailNow()
+	}
 	assert.Equal(t, result.Nodes[14].String(),
 		"MethodDeclaration_testUnpackEntryFromStreamToFile_"+newfileName)
 	assert.Equal(t, result.Counters[14], map[int]int{14: 1, 13: 1})
@@ -199,7 +212,7 @@ func TestShotnessConsume(t *testing.T) {
 		Change: &object.Change{
 			From: object.ChangeEntry{Name: newfileName},
 			To:   object.ChangeEntry{}},
-		Before: loadUast("uast2.pb"), After: nil,
+		Before: loadUast(t, "uast2.pb"), After: nil,
 	}
 	iresult, err = sh.Consume(state)
 	assert.Nil(t, err)