From 2dc5fd598244fe8fdf03152148bf7aaba5323433 Mon Sep 17 00:00:00 2001
From: TrueFurby <truefurby@gmail.com>
Date: Mon, 23 Jan 2017 04:05:14 +0100
Subject: [PATCH] Add support for multiple limit path prefixes

---
 README.md | 10 +++++-----
 main.go   | 32 ++++++++++++++++++++------------
 output.go | 33 +++++++++++++++++++++++++++++----
 3 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 272b597..07ac3f8 100644
--- a/README.md
+++ b/README.md
@@ -20,8 +20,8 @@ Purpose of this tool is to provide a visual overview of your program by using th
 
 - focus specific package in a program
 - group functions by package and methods by type
-- limit packages to custom prefix path
-- ignore packages containing custom prefix
+- limit packages to custom path prefixes
+- ignore packages containing path prefixes
 
 ### Output preview
 
@@ -100,11 +100,11 @@ Represents   | Style
 -focus string
       Focus package with import path or name. (default: main)
 -limit string
-      Limit package path to prefix.
+      Limit package paths to prefix. (separate multiple by comma)
 -group string
-      Grouping functions by [pkg, type] (separate multiple by comma).
+      Grouping functions by [pkg, type]. (separate multiple by comma)
 -ignore string
-      Ignore package paths with prefix (separate multiple by comma).
+      Ignore package paths with prefix. (separate multiple by comma)
 -minlen uint
       Minimum edge length (for wider output). (default: 2)
 -nodesep float
diff --git a/main.go b/main.go
index e7c8f0f..e53a01e 100644
--- a/main.go
+++ b/main.go
@@ -21,7 +21,7 @@ var Version = "0.0.0-src"
 
 var (
 	focusFlag   = flag.String("focus", "main", "Focus package with name or import path.")
-	limitFlag   = flag.String("limit", "", "Limit package path to prefix.")
+	limitFlag   = flag.String("limit", "", "Limit package paths to prefix. (separate multiple by comma)")
 	groupFlag   = flag.String("group", "", "Grouping functions by [pkg, type] (separate multiple by comma).")
 	ignoreFlag  = flag.String("ignore", "", "Ignore package paths with prefix (separate multiple by comma).")
 	testFlag    = flag.Bool("tests", false, "Include test code.")
@@ -44,14 +44,6 @@ func main() {
 		log.SetFlags(log.Lmicroseconds)
 	}
 
-	ignorePaths := []string{}
-	for _, p := range strings.Split(*ignoreFlag, ",") {
-		p = strings.TrimSpace(p)
-		if p != "" {
-			ignorePaths = append(ignorePaths, p)
-		}
-	}
-
 	groupBy := make(map[string]bool)
 	for _, g := range strings.Split(*groupFlag, ",") {
 		g := strings.TrimSpace(g)
@@ -65,13 +57,29 @@ func main() {
 		groupBy[g] = true
 	}
 
-	if err := run(&build.Default, *focusFlag, *limitFlag, groupBy, ignorePaths, *testFlag, flag.Args()); err != nil {
+	limitPaths := []string{}
+	for _, p := range strings.Split(*limitFlag, ",") {
+		p = strings.TrimSpace(p)
+		if p != "" {
+			limitPaths = append(limitPaths, p)
+		}
+	}
+
+	ignorePaths := []string{}
+	for _, p := range strings.Split(*ignoreFlag, ",") {
+		p = strings.TrimSpace(p)
+		if p != "" {
+			ignorePaths = append(ignorePaths, p)
+		}
+	}
+
+	if err := run(&build.Default, *focusFlag, groupBy, limitPaths, ignorePaths, *testFlag, flag.Args()); err != nil {
 		fmt.Fprintf(os.Stderr, "go-callvis: %s\n", err)
 		os.Exit(1)
 	}
 }
 
-func run(ctxt *build.Context, focus, limitPath string, groupBy map[string]bool, ignorePaths []string, tests bool, args []string) error {
+func run(ctxt *build.Context, focus string, groupBy map[string]bool, limitPaths, ignorePaths []string, tests bool, args []string) error {
 	if len(args) == 0 {
 		return fmt.Errorf("missing arguments")
 	}
@@ -154,7 +162,7 @@ func run(ctxt *build.Context, focus, limitPath string, groupBy map[string]bool,
 	logf("analysis took: %v", time.Since(t0))
 
 	return printOutput(mains[0].Pkg, result.CallGraph,
-		focusPkg, limitPath, ignorePaths, groupBy)
+		focusPkg, limitPaths, ignorePaths, groupBy)
 }
 
 func logf(f string, a ...interface{}) {
diff --git a/output.go b/output.go
index 6d0acfa..98e508d 100644
--- a/output.go
+++ b/output.go
@@ -14,7 +14,7 @@ import (
 
 var output io.Writer = os.Stdout
 
-func printOutput(mainPkg *types.Package, cg *callgraph.Graph, focusPkg *build.Package, limitPath string, ignorePaths []string, groupBy map[string]bool) error {
+func printOutput(mainPkg *types.Package, cg *callgraph.Graph, focusPkg *build.Package, limitPaths, ignorePaths []string, groupBy map[string]bool) error {
 	groupType := groupBy["type"]
 	groupPkg := groupBy["pkg"]
 
@@ -39,6 +39,32 @@ func printOutput(mainPkg *types.Package, cg *callgraph.Graph, focusPkg *build.Pa
 
 	cg.DeleteSyntheticNodes()
 
+	logf("%d limit prefixes: %v", len(limitPaths), limitPaths)
+	logf("%d ignore prefixes: %v", len(ignorePaths), ignorePaths)
+
+	var inLimit = func(from, to *callgraph.Node) bool {
+		if len(limitPaths) == 0 {
+			return true
+		}
+		var fromOk, toOk bool
+		fromPath := from.Func.Pkg.Pkg.Path()
+		toPath := to.Func.Pkg.Pkg.Path()
+		for _, p := range limitPaths {
+			if strings.HasPrefix(fromPath, p) {
+				fromOk = true
+			}
+			if strings.HasPrefix(toPath, p) {
+				toOk = true
+			}
+			if fromOk && toOk {
+				logf("in limit: %s -> %s", from, to)
+				return true
+			}
+		}
+		logf("NOT in limit: %s -> %s", from, to)
+		return false
+	}
+
 	err := callgraph.GraphVisitEdges(cg, func(edge *callgraph.Edge) error {
 		caller := edge.Caller
 		callee := edge.Callee
@@ -57,9 +83,8 @@ func printOutput(mainPkg *types.Package, cg *callgraph.Graph, focusPkg *build.Pa
 			return nil
 		}
 
-		// limit to path prefix
-		if !(strings.HasPrefix(callerPkg.Path(), limitPath) &&
-			strings.HasPrefix(calleePkg.Path(), limitPath)) {
+		// limit path prefixes
+		if !inLimit(caller, callee) {
 			return nil
 		}