From ebb532f365f7840d7d5e1e225508c63bf50e9a42 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Fri, 20 Apr 2018 09:26:32 +0200
Subject: [PATCH] jetCheck: don't retain generators to avoid OOME by capturing
 user objects

---
 .../src/org/jetbrains/jetCheck/NodeId.java    |  4 ++--
 .../org/jetbrains/jetCheck/StructureNode.java | 12 +++++------
 .../jetCheck/StatefulGeneratorTest.java       | 21 +++++++++++++++++++
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/jetCheck/src/org/jetbrains/jetCheck/NodeId.java b/jetCheck/src/org/jetbrains/jetCheck/NodeId.java
index c90fdfe81c8b..ac39bcb8378f 100644
--- a/jetCheck/src/org/jetbrains/jetCheck/NodeId.java
+++ b/jetCheck/src/org/jetbrains/jetCheck/NodeId.java
@@ -12,7 +12,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 class NodeId {
   private final AtomicInteger counter;
   final int number;
-  @Nullable final Generator<?> generator;
+  @Nullable final Integer generatorHash;
 
   NodeId(@NotNull Generator<?> generator) {
     this(new AtomicInteger(), generator);
@@ -20,7 +20,7 @@ class NodeId {
 
   private NodeId(AtomicInteger counter, @Nullable Generator<?> generator) {
     this.counter = counter;
-    this.generator = generator;
+    this.generatorHash = generator == null ? null : generator.getGeneratorFunction().hashCode();
     number = counter.getAndIncrement();
   }
 
diff --git a/jetCheck/src/org/jetbrains/jetCheck/StructureNode.java b/jetCheck/src/org/jetbrains/jetCheck/StructureNode.java
index a19927ae3ecd..13787836caf3 100644
--- a/jetCheck/src/org/jetbrains/jetCheck/StructureNode.java
+++ b/jetCheck/src/org/jetbrains/jetCheck/StructureNode.java
@@ -137,14 +137,14 @@ class StructureNode extends StructureElement {
     return false;
   }
 
-  private void findChildrenWithGenerator(@NotNull Generator<?> generator, List<StructureNode> result) {
+  private void findChildrenWithGenerator(int generatorHash, List<StructureNode> result) {
     for (StructureElement child : children) {
       if (child instanceof StructureNode) {
-        Generator<?> childGen = child.id.generator;
-        if (childGen != null && generator.getGeneratorFunction().equals(childGen.getGeneratorFunction())) {
+        Integer childGen = child.id.generatorHash;
+        if (childGen != null && generatorHash == childGen.intValue()) {
           result.add((StructureNode)child);
         } else {
-          ((StructureNode)child).findChildrenWithGenerator(generator, result);
+          ((StructureNode)child).findChildrenWithGenerator(generatorHash, result);
         }
       }
     }
@@ -152,9 +152,9 @@ class StructureNode extends StructureElement {
 
   @Nullable
   private ShrinkStep shrinkRecursion() {
-    if (id.generator != null) {
+    if (id.generatorHash != null) {
       List<StructureNode> sameGeneratorChildren = new ArrayList<>();
-      findChildrenWithGenerator(id.generator, sameGeneratorChildren);
+      findChildrenWithGenerator(id.generatorHash, sameGeneratorChildren);
       return tryReplacing(sameGeneratorChildren, 0);
     }
     
diff --git a/jetCheck/test/org/jetbrains/jetCheck/StatefulGeneratorTest.java b/jetCheck/test/org/jetbrains/jetCheck/StatefulGeneratorTest.java
index e581f710b828..05bc2734bb27 100644
--- a/jetCheck/test/org/jetbrains/jetCheck/StatefulGeneratorTest.java
+++ b/jetCheck/test/org/jetbrains/jetCheck/StatefulGeneratorTest.java
@@ -5,6 +5,7 @@ package org.jetbrains.jetCheck;
 
 import org.jetbrains.annotations.NotNull;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -131,6 +132,26 @@ public class StatefulGeneratorTest extends PropertyCheckerTestCase {
       sb.delete(start, end);
     };
   }
+
+  private ImperativeCommand heavyCommand() {
+    Object[] heavyObject = new Object[100_000];
+    heavyObject[42] = new Object();
+    return new ImperativeCommand() {
+      @Override
+      public void performCommand(@NotNull Environment env) {}
+
+      @Override
+      public String toString() {
+        return super.toString() + Arrays.toString(heavyObject);
+      }
+    };
+  }
+
+  public void testDontFailByOutOfMemoryDueToLeakingObjectsPassedIntoGenerators() {
+    PropertyChecker.customized().checkScenarios(() -> env -> 
+      env.executeCommands(from(data -> data.generate(sampledFrom(heavyCommand(), heavyCommand(), heavyCommand())))));
+  }
+
 }
 
 class InsertChar {