From 1f53d04893ea9b2736e9b2277bbba3f47b88711a Mon Sep 17 00:00:00 2001
From: Denys Smirnov <denis.smirnov.91@gmail.com>
Date: Sat, 6 Jan 2018 15:57:49 +0200
Subject: [PATCH] iterator: calculate real size for LinksTo(Fixed)

---
 graph/iterator.go         |  5 +++--
 graph/iterator/fixed.go   |  5 +++++
 graph/iterator/linksto.go | 35 +++++++++++++++++++++++++++++------
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/graph/iterator.go b/graph/iterator.go
index 65c9606ae..24a0405fb 100644
--- a/graph/iterator.go
+++ b/graph/iterator.go
@@ -175,13 +175,13 @@ type Iterator interface {
 
 // DescribeIterator returns a description of the iterator tree.
 func DescribeIterator(it Iterator) Description {
-	sz, _ := it.Size()
+	sz, exact := it.Size()
 	d := Description{
 		UID:  it.UID(),
 		Name: it.String(),
 		Type: it.Type(),
 		Tags: it.Tagger().Tags(),
-		Size: sz,
+		Size: sz, Exact: exact,
 	}
 	if sub := it.SubIterators(); len(sub) != 0 {
 		d.Iterators = make([]Description, 0, len(sub))
@@ -198,6 +198,7 @@ type Description struct {
 	Type      Type          `json:",omitempty"`
 	Tags      []string      `json:",omitempty"`
 	Size      int64         `json:",omitempty"`
+	Exact     bool          `json:",omitempty"`
 	Iterators []Description `json:",omitempty"`
 }
 
diff --git a/graph/iterator/fixed.go b/graph/iterator/fixed.go
index bf7390179..962e8b7e6 100644
--- a/graph/iterator/fixed.go
+++ b/graph/iterator/fixed.go
@@ -85,6 +85,11 @@ func (it *Fixed) Add(v graph.Value) {
 	it.values = append(it.values, v)
 }
 
+// Values returns a list of values stored in iterator. Slice should not be modified.
+func (it *Fixed) Values() []graph.Value {
+	return it.values
+}
+
 func (it *Fixed) String() string {
 	return fmt.Sprintf("Fixed(%v)", it.values)
 }
diff --git a/graph/iterator/linksto.go b/graph/iterator/linksto.go
index b89514fbf..21f90c110 100644
--- a/graph/iterator/linksto.go
+++ b/graph/iterator/linksto.go
@@ -85,6 +85,7 @@ func (it *LinksTo) Tagger() *graph.Tagger {
 func (it *LinksTo) Clone() graph.Iterator {
 	out := NewLinksTo(it.qs, it.primaryIt.Clone(), it.dir)
 	out.tags.CopyFrom(it)
+	out.runstats.Size, out.runstats.ExactSize = it.runstats.Size, it.runstats.ExactSize
 	return out
 }
 
@@ -211,21 +212,43 @@ func (it *LinksTo) Type() graph.Type { return graph.LinksTo }
 func (it *LinksTo) Stats() graph.IteratorStats {
 	subitStats := it.primaryIt.Stats()
 	// TODO(barakmich): These should really come from the quadstore itself
-	fanoutFactor := int64(20)
 	checkConstant := int64(1)
 	nextConstant := int64(2)
-	return graph.IteratorStats{
+	st := graph.IteratorStats{
 		NextCost:     nextConstant + subitStats.NextCost,
 		ContainsCost: checkConstant + subitStats.ContainsCost,
-		Size:         fanoutFactor * subitStats.Size,
-		ExactSize:    false,
 		Next:         it.runstats.Next,
 		Contains:     it.runstats.Contains,
 		ContainsNext: it.runstats.ContainsNext,
 	}
+	st.Size, st.ExactSize = it.Size()
+	return st
 }
 
 func (it *LinksTo) Size() (int64, bool) {
-	st := it.Stats()
-	return st.Size, st.ExactSize
+	if it.runstats.Size != 0 {
+		return it.runstats.Size, it.runstats.ExactSize
+	}
+	if fixed, ok := it.primaryIt.(*Fixed); ok {
+		// get real sizes from sub iterators
+		var (
+			sz    int64
+			exact = true
+		)
+		for _, v := range fixed.Values() {
+			sit := it.qs.QuadIterator(it.dir, v)
+			n, ex := sit.Size()
+			sit.Close()
+			sz += n
+			exact = exact && ex
+		}
+		it.runstats.Size, it.runstats.ExactSize = sz, exact
+		return sz, exact
+	}
+	// TODO(barakmich): It should really come from the quadstore itself
+	const fanoutFactor = 20
+	sz, _ := it.primaryIt.Size()
+	sz *= fanoutFactor
+	it.runstats.Size, it.runstats.ExactSize = sz, false
+	return sz, false
 }