feat(#16): ref points in edge attrs

williamfzc · Jul 16, 2023 · 67f41b4 · 67f41b4
1 parent d087bdf
commit 67f41b4
Show file tree

Hide file tree

Showing 8 changed files with 102 additions and 10 deletions.
diff --git a/example/api_test.go b/example/api_test.go
@@ -29,7 +29,7 @@ func TestFunc(t *testing.T) {
 	})
 
 	t.Run("func", func(t *testing.T) {
-		functions := funcGraph.GetFunctionsByFile("graph/function/api_query_func.go")
+		functions := funcGraph.GetFunctionsByFile("graph/function/api_query.go")
 		for _, each := range functions {
 			// about this function
 			log.Infof("func: %v", each.Id())

diff --git a/graph/common/edge.go b/graph/common/edge.go
@@ -0,0 +1,9 @@
+package common
+
+type EdgeStorage struct {
+	RefLines map[int]struct{}
+}
+
+func NewEdgeStorage() *EdgeStorage {
+	return &EdgeStorage{RefLines: make(map[int]struct{})}
+}
diff --git a/graph/file/api_query.go b/graph/file/api_query.go
@@ -1,6 +1,11 @@
 package file
 
-import "github.com/sirupsen/logrus"
+import (
+	"fmt"
+
+	"github.com/sirupsen/logrus"
+	"github.com/williamfzc/srctx/graph/common"
+)
 
 func (fg *Graph) GetById(id string) *Vertex {
 	v, err := fg.G.Vertex(id)
@@ -35,3 +40,14 @@ func (fg *Graph) ListEntries() []*Vertex {
 		return vertex.ContainTag(TagEntry)
 	})
 }
+
+func (fg *Graph) RelationBetween(a string, b string) (*common.EdgeStorage, error) {
+	edge, err := fg.G.Edge(a, b)
+	if err != nil {
+		return nil, err
+	}
+	if ret, ok := edge.Properties.Data.(*common.EdgeStorage); ok {
+		return ret, nil
+	}
+	return nil, fmt.Errorf("failed to convert %v", edge.Properties.Data)
+}
diff --git a/graph/file/api_test.go b/graph/file/api_test.go
@@ -5,9 +5,9 @@ import (
 	"runtime"
 	"testing"
 
-	"github.com/williamfzc/srctx/graph/file"
-
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
+	"github.com/williamfzc/srctx/graph/file"
 )
 
 func TestFileGraph(t *testing.T) {
@@ -41,4 +41,11 @@ func TestFileGraph(t *testing.T) {
 		err = fileGraph.DrawG6Html("b.html")
 		assert.Nil(t, err)
 	})
+
+	t.Run("Relation", func(t *testing.T) {
+		edgeStorage, err := fileGraph.RelationBetween("graph/function/api_query.go", "graph/function/api_query_test.go")
+		assert.Nil(t, err)
+		log.Debugf("ref lines: %v", edgeStorage.RefLines)
+		assert.NotEmpty(t, edgeStorage.RefLines)
+	})
 }
diff --git a/graph/file/graph.go b/graph/file/graph.go
@@ -3,6 +3,7 @@ package file
 import (
 	"github.com/dominikbraun/graph"
 	log "github.com/sirupsen/logrus"
+	"github.com/williamfzc/srctx/graph/common"
 	"github.com/williamfzc/srctx/object"
 	"github.com/williamfzc/srctx/parser"
 )
@@ -150,8 +151,21 @@ func CreateFileGraph(relationship *object.SourceContext) (*Graph, error) {
 				if eachSrcFile == targetFile {
 					continue
 				}
-				_ = g.G.AddEdge(eachSrcFile, targetFile)
-				_ = g.Rg.AddEdge(targetFile, eachSrcFile)
+
+				refLineNumber := eachRef.LineNumber()
+				if edge, err := g.G.Edge(eachSrcFile, targetFile); err == nil {
+					storage := edge.Properties.Data.(*common.EdgeStorage)
+					storage.RefLines[refLineNumber] = struct{}{}
+				} else {
+					_ = g.G.AddEdge(eachSrcFile, targetFile, graph.EdgeData(common.NewEdgeStorage()))
+				}
+
+				if edge, err := g.Rg.Edge(targetFile, eachSrcFile); err == nil {
+					storage := edge.Properties.Data.(*common.EdgeStorage)
+					storage.RefLines[refLineNumber] = struct{}{}
+				} else {
+					_ = g.Rg.AddEdge(targetFile, eachSrcFile, graph.EdgeData(common.NewEdgeStorage()))
+				}
 			}
 		}
 	}

diff --git a/graph/function/api_query.go b/graph/function/api_query.go
@@ -1,6 +1,10 @@
 package function
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/williamfzc/srctx/graph/common"
+)
 
 func (fg *Graph) GetFunctionsByFile(fileName string) []*Vertex {
 	if item, ok := fg.Cache[fileName]; ok {
@@ -57,3 +61,14 @@ func (fg *Graph) ListEntries() []*Vertex {
 		return funcVertex.ContainTag(TagEntry)
 	})
 }
+
+func (fg *Graph) RelationBetween(a string, b string) (*common.EdgeStorage, error) {
+	edge, err := fg.g.Edge(a, b)
+	if err != nil {
+		return nil, err
+	}
+	if ret, ok := edge.Properties.Data.(*common.EdgeStorage); ok {
+		return ret, nil
+	}
+	return nil, fmt.Errorf("failed to convert %v", edge.Properties.Data)
+}
diff --git a/graph/function/api_query_test.go b/graph/function/api_query_test.go
@@ -5,6 +5,8 @@ import (
 	"runtime"
 	"testing"
 
+	log "github.com/sirupsen/logrus"
+
 	"github.com/opensibyl/sibyl2/pkg/core"
 	"github.com/stretchr/testify/assert"
 )
@@ -44,4 +46,18 @@ func TestQuery(t *testing.T) {
 			}
 		}
 	})
+
+	t.Run("Relation", func(t *testing.T) {
+		testFuncs := fg.GetFunctionsByFile("graph/function/api_query_test.go")
+		assert.NotEmpty(t, testFuncs)
+
+		function, err := fg.GetById("graph/function/api_query.go:#9-#14:function|*Graph|GetFunctionsByFile|string|")
+		assert.Nil(t, err)
+		assert.NotNil(t, function)
+
+		edgeStorage, err := fg.RelationBetween(function.Id(), testFuncs[0].Id())
+		assert.Nil(t, err)
+		log.Debugf("ref lines: %v", edgeStorage.RefLines)
+		assert.NotEmpty(t, edgeStorage.RefLines)
+	})
 }
diff --git a/graph/function/graph.go b/graph/function/graph.go
@@ -5,6 +5,8 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/williamfzc/srctx/graph/common"
+
 	"github.com/dominikbraun/graph"
 	"github.com/opensibyl/sibyl2/pkg/core"
 	"github.com/opensibyl/sibyl2/pkg/extractor"
@@ -159,9 +161,22 @@ func CreateFuncGraph(fact *FactStorage, relationship *object.SourceContext) (*Gr
 					// eachPossibleFunc 's range contains eachFunc 's ref
 					// so eachPossibleFunc calls eachFunc
 					if eachPossibleFunc.GetSpan().ContainLine(eachRef.IndexLineNumber()) {
-						log.Debugf("%v refed in %s#%v", eachFunc.Id(), refFile, eachRef.LineNumber())
-						_ = fg.g.AddEdge(eachFunc.Id(), eachPossibleFunc.Id())
-						_ = fg.rg.AddEdge(eachPossibleFunc.Id(), eachFunc.Id())
+						refLineNumber := eachRef.LineNumber()
+						log.Debugf("%v refed in %s#%v", eachFunc.Id(), refFile, refLineNumber)
+
+						if edge, err := fg.g.Edge(eachFunc.Id(), eachPossibleFunc.Id()); err == nil {
+							storage := edge.Properties.Data.(*common.EdgeStorage)
+							storage.RefLines[refLineNumber] = struct{}{}
+						} else {
+							_ = fg.g.AddEdge(eachFunc.Id(), eachPossibleFunc.Id(), graph.EdgeData(common.NewEdgeStorage()))
+						}
+
+						if edge, err := fg.rg.Edge(eachPossibleFunc.Id(), eachFunc.Id()); err == nil {
+							storage := edge.Properties.Data.(*common.EdgeStorage)
+							storage.RefLines[refLineNumber] = struct{}{}
+						} else {
+							_ = fg.rg.AddEdge(eachPossibleFunc.Id(), eachFunc.Id(), graph.EdgeData(common.NewEdgeStorage()))
+						}
 					}
 				}
 			}