package water;
import water.util.ArrayUtils;
import water.util.Log;

import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Stack;

// A "scope" for tracking Key lifetimes.
//
// A Scope defines a *SINGLE THREADED* local lifetime management context,
// stored in Thread Local Storage.  Scopes can be explicitly entered or exited.
// User keys created by this thread are tracked, and deleted when the scope is
// exited.  Since enter & exit are explicit, failure to exit means the Keys
// leak (there is no reliable thread-on-exit cleanup action).  You must call
// Scope.exit() at some point.  Only user keys & Vec keys are tracked.
//
// Scopes support nesting.  Scopes support partial cleanup: you can list Keys
// you'd like to keep in the exit() call.  These will be "bumped up" to the
// higher nested scope - or escaped and become untracked at the top-level.

public class Scope {
  // Thread-based Key lifetime tracking
  static private final ThreadLocal<Scope> _scope = new ThreadLocal<Scope>() {
    @Override protected Scope initialValue() { return new Scope(); }
  };
  private final Stack<HashSet<Key>> _keys = new Stack<>();

  static public void enter() { _scope.get()._keys.push(new HashSet<Key>()); }
  static public void exit () { exit(null); }

  static public Key[] exit(Key... keep) {
    Key[] arrkeep = null;
    if (keep != null) arrkeep = (Key[]) Arrays.asList(keep).toArray();
    Arrays.sort(arrkeep);
    Stack<HashSet<Key>> keys = _scope.get()._keys;
    if (keys.size() > 0) {
      for (Key key : keys.pop()) {
        int found = -1;
        if (arrkeep!=null) found = Arrays.binarySearch(arrkeep, key);
        if (found >= 0) {
          Log.info("Not deleting key " + key);
        }
        if (arrkeep.length == 0 || found < 0) {
          Log.info("Deleting key " + key);
          Keyed.remove(key);
        }
      }
    }
    return keep;
  }

  static public void track( Key k ) {
    if( !k.user_allowed() && !k.isVec() ) return; // Not tracked
    Scope scope = _scope.get();                   // Pay the price of T.L.S. lookup
    if( scope == null ) return; // Not tracking this thread
    if( scope._keys.size() == 0 ) return; // Tracked in the past, but no scope now
    scope._keys.peek().add(k);            // Track key
  }
}

@@ -12,158 +12,27 @@
 er;%0A
-import water.util.ArrayUtils;%0Aimport water.util.Log;%0A%0Aimport java.util.Arrays;%0Aimport java.util.HashSet;%0Aimport java.util.List;%0Aimport java.util.Stack
+%0Aimport java.util.*
 ;%0A%0A/
@@ -1211,28 +1211,45 @@
     
-Key%5B%5D arrkeep = null
+List%3CKey%3E keylist = new ArrayList%3C%3E()
 ;%0A  
@@ -1272,25 +1272,28 @@
 ll) 
-arrkeep =
+%7B%0A      for
  (Key
-%5B%5D)
+ k :
  Arr
@@ -1308,16 +1308,113 @@
 st(keep)
+) %7B%0A        if (k != null)%0A          keylist.add(k);%0A      %7D%0A    %7D%0A    Object%5B%5D arrkeep = keylist
 .toArray
@@ -1599,18 +1599,20 @@
 (arrkeep
+ 
 !=
+ 
 null) fo
@@ -1648,24 +1648,26 @@
 keep, key);%0A
+//
         if (
@@ -1676,24 +1676,26 @@
 und %3E= 0) %7B%0A
+//
           Lo
@@ -1721,32 +1721,34 @@
 g key %22 + key);%0A
+//
         %7D%0A      
@@ -1785,24 +1785,26 @@
 ound %3C 0) %7B%0A
+//
           Lo
@@ -2239,16 +2239,101 @@
 ope now%0A
+    if (!scope._keys.peek().contains(k)) %7B%0A//      Log.info(%22Tracking key: %22 + k);%0A  
     scop
@@ -2379,14 +2379,20 @@
 ack key%0A
+    %7D%0A
   %7D%0A%7D%0A

Improve Scope code for exit(Key... key), allow null to be in varargs.