Add server side filtering of treemap data.

Author: baffles

codepulse/src/main/scala/com/secdec/codepulse/data/trace/TreeBuilder.scala

@@ -14,6 +14,7 @@ case class PackageTreeNode(id: Option[Int], kind: CodeTreeNodeKind, label: Strin
   * @author robertf
   */
 class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
+	/** Full set of tree roots and nodes */
 	lazy val (roots, nodes) = {
 		val roots = List.newBuilder[Int]
 		val nodes = Map.newBuilder[Int, TreeNodeData]
@@ -39,6 +40,7 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 		(roots.result.map(buildNode), nodeMap)
 	}
 
+	/** Full package tree, with self nodes */
 	lazy val packageTree = {
 		// build up a package tree with the relevant data
 
@@ -85,4 +87,48 @@ class TreeBuilder(treeNodeData: TreeNodeDataAccess) {
 
 		roots.map(transform)
 	}
+
+	/** Projects a tree containing the selected packages and their immediate children */
+	def projectTree(selectedNodes: Set[Int]) = {
+		val incidentalNodes = collection.mutable.HashSet.empty[Int]
+
+		// recursively mark all parents of `selectedNodes` as incidental nodes (partially accepted)
+		def markIncidentalPath(node: Int) {
+			incidentalNodes += node
+
+			for {
+				node <- nodes get node
+				parent <- node.parentId
+			} markIncidentalPath(parent)
+		}
+
+		selectedNodes.foreach(markIncidentalPath)
+
+		// build the projected tree
+		def filterNode(node: TreeNode): Option[TreeNode] = {
+			def isSubstantialChild(node: TreeNode) = node.data.kind != CodeTreeNodeKind.Grp && node.data.kind != CodeTreeNodeKind.Pkg
+
+			// only include this node if it is incidental
+			if (incidentalNodes contains node.data.id) {
+				val isSelected = selectedNodes contains node.data.id
+
+				// filter children; don't include substantial data if only incidental
+				// the only exception to this is if that child was requested specifically
+				val filteredChildren = node.children.flatMap {
+					case child if isSubstantialChild(child) =>
+						if (isSelected) Some(child)
+						else filterNode(child)
+
+					case child => filterNode(child)
+				}
+
+				if (isSelected || !filteredChildren.isEmpty)
+					Some(TreeNode(node.data, filteredChildren))
+				else
+					None
+			} else None
+		}
+
+		roots.flatMap(filterNode)
+	}
 }

codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceAPIServer.scala

@@ -26,13 +26,11 @@ import scala.concurrent.duration.DurationInt
 import scala.language.implicitConversions
 import scala.util.Failure
 import scala.util.Success
-
 import org.joda.time.format.DateTimeFormat
-
 import com.secdec.codepulse.data.trace._
 import com.secdec.codepulse.pages.traces.TraceDetailsPage
-
 import akka.actor.Cancellable
+import net.liftweb.common.Full
 import net.liftweb.common.Loggable
 import net.liftweb.http._
 import net.liftweb.http.rest._
@@ -293,9 +291,16 @@ class TraceAPIServer(manager: TraceManager) extends RestHelper with Loggable {
 		case Paths.PackageTree(target) Get req =>
 			PackageTreeStreamer.streamPackageTree(target.treeBuilder.packageTree)
 
-		// GET the trace's treemap data as json
-		case Paths.Treemap(target) Get req =>
-			TreemapDataStreamer.streamTreemapData(target.traceData)
+		// stream a projected treemap based on the POSTed selected packages
+		case Paths.Treemap(target) Post req =>
+			req.param("nodes") match {
+				case Full(packages) =>
+					val ids = packages.split(',').flatMap(AsInt.unapply).toSet
+					val tree = target.treeBuilder.projectTree(ids)
+					TreemapDataStreamer.streamTreemapData(tree)
+
+				case _ => BadResponse()
+			}
 
 		case Paths.TreeInstrumentation(target) Put req =>
 			def getBool(j: JValue) = j match {

codepulse/src/main/scala/com/secdec/codepulse/tracer/TreemapDataStreamer.scala

@@ -34,27 +34,33 @@ import net.liftweb.json.Printer
 object TreemapDataStreamer {
 	private val Json = new JsonFactory
 
-	private def writeJson(jg: JsonGenerator)(node: TreeNodeData) {
-		jg writeFieldName node.id.toString
+	private def writeJson(jg: JsonGenerator)(node: TreeNode) {
 		jg.writeStartObject
 
-		for (parentId <- node.parentId) jg.writeNumberField("parentId", parentId)
-		jg.writeStringField("name", node.label)
-		jg.writeStringField("kind", node.kind.label)
-		for (size <- node.size) jg.writeNumberField("lineCount", size)
-		for (traced <- node.traced) jg.writeBooleanField("traced", traced)
+		jg.writeNumberField("id", node.data.id)
+		for (parentId <- node.data.parentId) jg.writeNumberField("parentId", parentId)
+		jg.writeStringField("name", node.data.label)
+		jg.writeStringField("kind", node.data.kind.label)
+		for (size <- node.data.size) jg.writeNumberField("lineCount", size)
+		for (traced <- node.data.traced) jg.writeBooleanField("traced", traced)
+
+		if (!node.children.isEmpty) {
+			jg writeArrayFieldStart "children"
+			node.children.foreach(writeJson(jg))
+			jg.writeEndArray
+		}
 
 		jg.writeEndObject
 	}
 
-	def streamTreemapData(traceData: TraceData): OutputStreamResponse = {
+	def streamTreemapData(tree: List[TreeNode]): OutputStreamResponse = {
 		def writeData(out: OutputStream) {
 			val jg = Json createGenerator out
 
 			try {
-				jg.writeStartObject
-				traceData.treeNodeData.foreach(writeJson(jg))
-				jg.writeEndObject
+				jg.writeStartArray
+				tree.foreach(writeJson(jg))
+				jg.writeEndArray
 			} finally jg.close
 		}