LUCENE-5487: also wrap TopScorer in AssertingWeight

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/branches/lucene5487@1574169 13f79535-47bb-0310-9956-ffa450edef68

lucene/test-framework/src/java/org/apache/lucene/search/AssertingIndexSearcher.java

@@ -90,6 +90,8 @@ public class AssertingIndexSearcher extends IndexSearcher {
 
   @Override
   protected void search(List<AtomicReaderContext> leaves, Weight weight, Collector collector) throws IOException {
+    // nocommit shouldn't we
+    // AssertingCollector.wrap(collector) here?
     super.search(leaves, AssertingWeight.wrap(random, weight), collector);
   }
 

lucene/test-framework/src/java/org/apache/lucene/search/AssertingScorer.java

@@ -32,9 +32,6 @@ import org.apache.lucene.util.VirtualMethod;
 /** Wraps a Scorer with additional checks */
 public class AssertingScorer extends Scorer {
 
-  //private static final VirtualMethod<Scorer> SCORE_COLLECTOR = new VirtualMethod<Scorer>(Scorer.class, "score", Collector.class);
-  //private static final VirtualMethod<Scorer> SCORE_COLLECTOR_RANGE = new VirtualMethod<Scorer>(Scorer.class, "score", Collector.class, int.class, int.class);
-
   // we need to track scorers using a weak hash map because otherwise we
   // could loose references because of eg.
   // AssertingScorer.score(Collector) which needs to delegate to work correctly
@@ -100,41 +97,6 @@ public class AssertingScorer extends Scorer {
     return score;
   }
 
-  // nocommit make an AssertingTopScorer
-  /*
-  @Override
-  public void score(Collector collector) throws IOException {
-    if (SCORE_COLLECTOR.isOverriddenAsOf(this.in.getClass())) {
-      if (random.nextBoolean()) {
-        try {
-          final boolean remaining = in.score(collector, DocsEnum.NO_MORE_DOCS, in.nextDoc());
-          assert !remaining;
-        } catch (UnsupportedOperationException e) {
-          in.score(collector);
-        }
-      } else {
-        in.score(collector);
-      }
-    } else {
-      // score(Collector) has not been overridden, use the super method in
-      // order to benefit from all assertions
-      super.score(collector);
-    }
-  }
-
-  @Override
-  public boolean score(Collector collector, int max, int firstDocID) throws IOException {
-    assert topScorer != TopScorer.NO;
-    if (SCORE_COLLECTOR_RANGE.isOverriddenAsOf(this.in.getClass())) {
-      return in.score(collector, max, firstDocID);
-    } else {
-      // score(Collector,int,int) has not been overridden, use the super
-      // method in order to benefit from all assertions
-      return super.score(collector, max, firstDocID);
-    }
-  }
-  */
-
   @Override
   public Collection<ChildScorer> getChildren() {
     // We cannot hide that we hold a single child, else

lucene/test-framework/src/java/org/apache/lucene/search/AssertingTopScorer.java

@@ -0,0 +1,90 @@
+package org.apache.lucene.search;
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import java.io.IOException;
+import java.lang.ref.WeakReference;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Random;
+import java.util.WeakHashMap;
+
+import org.apache.lucene.index.DocsEnum;
+import org.apache.lucene.util.VirtualMethod;
+
+/** Wraps a Scorer with additional checks */
+public class AssertingTopScorer extends TopScorer {
+
+  private static final VirtualMethod<TopScorer> SCORE_COLLECTOR = new VirtualMethod<TopScorer>(TopScorer.class, "score", Collector.class);
+  private static final VirtualMethod<TopScorer> SCORE_COLLECTOR_RANGE = new VirtualMethod<TopScorer>(TopScorer.class, "score", Collector.class, int.class);
+
+  // we need to track scorers using a weak hash map because otherwise we
+  // could loose references because of eg.
+  // AssertingScorer.score(Collector) which needs to delegate to work correctly
+  private static Map<TopScorer, WeakReference<AssertingTopScorer>> ASSERTING_INSTANCES = Collections.synchronizedMap(new WeakHashMap<TopScorer, WeakReference<AssertingTopScorer>>());
+
+  public static TopScorer wrap(Random random, TopScorer other) {
+    if (other == null || other instanceof AssertingTopScorer) {
+      return other;
+    }
+    final AssertingTopScorer assertScorer = new AssertingTopScorer(random, other);
+    ASSERTING_INSTANCES.put(other, new WeakReference<AssertingTopScorer>(assertScorer));
+    return assertScorer;
+  }
+
+  public static boolean shouldWrap(TopScorer inScorer) {
+    return SCORE_COLLECTOR.isOverriddenAsOf(inScorer.getClass()) || SCORE_COLLECTOR_RANGE.isOverriddenAsOf(inScorer.getClass());
+  }
+
+  final Random random;
+  final TopScorer in;
+
+  private AssertingTopScorer(Random random, TopScorer in) {
+    this.random = random;
+    this.in = in;
+  }
+
+  public TopScorer getIn() {
+    return in;
+  }
+
+  @Override
+  public void score(Collector collector) throws IOException {
+    if (random.nextBoolean()) {
+      try {
+        final boolean remaining = in.score(collector, DocsEnum.NO_MORE_DOCS);
+        assert !remaining;
+      } catch (UnsupportedOperationException e) {
+        in.score(collector);
+      }
+    } else {
+      in.score(collector);
+    }
+  }
+
+  @Override
+  public boolean score(Collector collector, int max) throws IOException {
+    return in.score(collector, max);
+  }
+
+  @Override
+  public String toString() {
+    return "AssertingTopScorer(" + in + ")";
+  }
+}

lucene/test-framework/src/java/org/apache/lucene/search/AssertingWeight.java

@@ -69,10 +69,17 @@ class AssertingWeight extends Weight {
   public TopScorer topScorer(AtomicReaderContext context, boolean scoreDocsInOrder, Bits acceptDocs) throws IOException {
     // if the caller asks for in-order scoring or if the weight does not support
     // out-of order scoring then collection will have to happen in-order.
-    // nocommit add wrapping:
     TopScorer inScorer = in.topScorer(context, scoreDocsInOrder, acceptDocs);
-    //return AssertingScorer.wrap(new Random(random.nextLong()), inScorer);
-    return inScorer;
+    if (inScorer == null) {
+      return null;
+    }
+    if (AssertingTopScorer.shouldWrap(inScorer)) {
+      return AssertingTopScorer.wrap(new Random(random.nextLong()), inScorer);
+    } else {
+      // Let super wrap this.scorer instead, so we use
+      // AssertingScorer:
+      return super.topScorer(context, scoreDocsInOrder, acceptDocs);
+    }
   }
 
   @Override