Try to determine EndPosition if it does not exist

README.md

@@ -183,16 +183,17 @@ can be visualized with t-SNE implemented in TF Projector.
 #### Structural hotness
 
 ```
+      15  jinja2/environment.py:compile [FunctionDef]
+       9  jinja2/environment.py:parse [FunctionDef]
+       8  jinja2/runtime.py:call [FunctionDef]
        8  jinja2/runtime.py:__init__ [FunctionDef]
+       8  jinja2/runtime.py:__call__ [FunctionDef]
        7  jinja2/ext.py:parse [FunctionDef]
+       7  jinja2/runtime.py:length [FunctionDef]
+       7  jinja2/environment.py:get_template [FunctionDef]
+       7  jinja2/runtime.py:_fail_with_undefined_error [FunctionDef]
        6  jinja2/optimizer.py:optimize [FunctionDef]
-       6  jinja2/compiler.py:inspect [FunctionDef]
-       6  jinja2/loaders.py:__init__ [FunctionDef]
-       6  jinja2/exceptions.py:__init__ [FunctionDef]
-       6  jinja2/runtime.py:__getitem__ [FunctionDef]
-       5  jinja2/optimizer.py:visit_For [FunctionDef]
-       5  jinja2/parser.py:parse_tuple [FunctionDef]
-       5  jinja2/environment.py:compile [FunctionDef]
+
 ```
 
 Thanks to Babelfish, hercules is able to measure how many times each structural unit has been modified.

diff_refiner_test.go

@@ -92,9 +92,9 @@ func TestFileDiffRefinerConsume(t *testing.T) {
 	assert.Equal(t, len(oldDiff.Diffs)+1, len(newDiff.Diffs))
 	assert.Equal(t, dmp.DiffText2(oldDiff.Diffs), dmp.DiffText2(newDiff.Diffs))
 	// Some hardcoded length checks
-	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[5].Text), 11)
-	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[6].Text), 41)
-	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[7].Text), 231)
+	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[6].Text), 11)
+	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[7].Text), 41)
+	assert.Equal(t, utf8.RuneCountInString(newDiff.Diffs[8].Text), 231)
 }
 
 func TestFileDiffRefinerConsumeNoUast(t *testing.T) {

shotness.go

@@ -221,15 +221,34 @@ func (shotness *ShotnessAnalysis) Consume(deps map[string]interface{}) (map[stri
 		genLine2Node := func(nodes map[string]*uast.Node, linesNum int) [][]*uast.Node {
 			res := make([][]*uast.Node, linesNum)
 			for _, node := range nodes {
-				if node.StartPosition != nil && node.EndPosition != nil {
-					for l := node.StartPosition.Line; l <= node.EndPosition.Line; l++ {
-						lineNodes := res[l-1]
-						if lineNodes == nil {
-							lineNodes = []*uast.Node{}
+				if node.StartPosition == nil {
+					continue
+				}
+				startLine := node.StartPosition.Line
+				endLine := node.StartPosition.Line
+				if node.EndPosition != nil {
+					endLine = node.EndPosition.Line
+				} else {
+					// we need to determine node.EndPosition.Line
+					VisitEachNode(node, func(child *uast.Node) {
+						if child.StartPosition != nil {
+							candidate := child.StartPosition.Line
+							if child.EndPosition != nil {
+								candidate = child.EndPosition.Line
+							}
+							if candidate > endLine {
+								endLine = candidate
+							}
 						}
-						lineNodes = append(lineNodes, node)
-						res[l-1] = lineNodes
+					})
+				}
+				for l := startLine; l <= endLine; l++ {
+					lineNodes := res[l-1]
+					if lineNodes == nil {
+						lineNodes = []*uast.Node{}
 					}
+					lineNodes = append(lineNodes, node)
+					res[l-1] = lineNodes
 				}
 			}
 			return res

shotness_test.go

@@ -57,7 +57,7 @@ func TestShotnessRegistration(t *testing.T) {
 	assert.Equal(t, tp.Elem().Name(), "ShotnessAnalysis")
 }
 
-func bakeShotness(t *testing.T) (*ShotnessAnalysis, ShotnessResult) {
+func bakeShotness(t *testing.T, eraseEndPosition bool) (*ShotnessAnalysis, ShotnessResult) {
 	sh := fixtureShotness()
 	bytes1, err := ioutil.ReadFile(path.Join("test_data", "1.java"))
 	assert.Nil(t, err)
@@ -81,6 +81,11 @@ func bakeShotness(t *testing.T) (*ShotnessAnalysis, ShotnessResult) {
 		assert.Nil(t, err)
 		node := uast.Node{}
 		proto.Unmarshal(bytes, &node)
+		if eraseEndPosition {
+			VisitEachNode(&node, func(child *uast.Node) {
+				child.EndPosition = nil
+			})
+		}
 		return &node
 	}
 	state[DependencyUastChanges] = uastChanges
@@ -193,8 +198,14 @@ func TestShotnessConsume(t *testing.T) {
 	assert.Len(t, sh.files, 0)
 }
 
+func TestShotnessConsumeNoEnd(t *testing.T) {
+	_, result1 := bakeShotness(t, false)
+	_, result2 := bakeShotness(t, true)
+	assert.Equal(t, result1, result2)
+}
+
 func TestShotnessSerializeText(t *testing.T) {
-	sh, result := bakeShotness(t)
+	sh, result := bakeShotness(t, false)
 	buffer := &bytes.Buffer{}
 	sh.Serialize(result, false, buffer)
 	assert.Equal(t, buffer.String(), `  - name: testAddEntry
@@ -291,7 +302,7 @@ func TestShotnessSerializeText(t *testing.T) {
 }
 
 func TestShotnessSerializeBinary(t *testing.T) {
-	sh, result := bakeShotness(t)
+	sh, result := bakeShotness(t, false)
 	buffer := &bytes.Buffer{}
 	sh.Serialize(result, true, buffer)
 	message := pb.ShotnessAnalysisResults{}

test_data/2.java

@@ -14,7 +14,6 @@
  *    limitations under the License.
  */
 import java.io.BufferedInputStream;
-import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.FileReader;