package water;
import java.util.Arrays;
import java.util.HashSet;
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 () {
    Key k = null;
    exit(k);
  }

  static public Key exit(Key keep) {
    Stack<HashSet<Key>> keys = _scope.get()._keys;
    if (keys.size() > 0) {
      for (Key key : keys.pop()) {
        if (keep == null || !Arrays.equals(key._kb, keep._kb)) {
          Keyed.remove(key);
        }
      }
    }
    return keep;
  }

  static public Key[] exit(Key... keep) { throw H2O.unimpl(); }

  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
  }
}

@@ -8,16 +8,70 @@
  water;%0A
+import water.util.ArrayUtils;%0Aimport water.util.Log;%0A%0A
 import j
@@ -113,16 +113,39 @@
 ashSet;%0A
+import java.util.List;%0A
 import j
@@ -1280,41 +1280,20 @@
 () %7B
-%0A    Key k = null;%0A    exit(k);%0A 
+ exit(null);
  %7D%0A%0A
@@ -1311,16 +1311,18 @@
 blic Key
+%5B%5D
  exit(Ke
@@ -1322,16 +1322,19 @@
 exit(Key
+...
  keep) %7B
@@ -1342,181 +1342,507 @@
     
-Stack%3CHashSet%3CKey%3E%3E keys = _scope.get()._keys;%0A    if (keys.size() %3E 0) %7B%0A      for (Key key : keys.pop()) %7B%0A        if (keep == null %7C%7C !Arrays.equals(key._kb, keep._kb)) %7B
+Key%5B%5D arrkeep = null;%0A    if (keep != null) arrkeep = (Key%5B%5D) Arrays.asList(keep).toArray();%0A    Arrays.sort(arrkeep);%0A    Stack%3CHashSet%3CKey%3E%3E keys = _scope.get()._keys;%0A    if (keys.size() %3E 0) %7B%0A      for (Key key : keys.pop()) %7B%0A        int found = -1;%0A        if (arrkeep!=null) found = Arrays.binarySearch(arrkeep, key);%0A        if (found %3E= 0) %7B%0A          Log.info(%22Not deleting key %22 + key);%0A        %7D%0A        if (arrkeep.length == 0 %7C%7C found %3C 0) %7B%0A          Log.info(%22Deleting key %22 + key);
 %0A   
@@ -1917,73 +1917,8 @@
  %7D%0A%0A
-  static public Key%5B%5D exit(Key... keep) %7B throw H2O.unimpl(); %7D%0A%0A
   st

Implement exit(Key... keep) method for scope.