StateScope.java

package de.bitbrain.guice;

import java.util.HashMap;
import java.util.Map;

import com.google.inject.Key;
import com.google.inject.Provider;
import com.google.inject.Scope;

public class StateScope implements Scope {
  
  private Class<?> state;
  
  private final ThreadLocal<Map<Class<?>, Object>> values
  = new ThreadLocal<Map<Class<?>, Object>>();
  private final Map<Key<?>, Class<?>> keys = new HashMap<Key<?>, Class<?>>();
  
  public void enter(Class<?> state) {
    this.state = state;
  }
  
  public void leave() {
    values.remove();
  }

  public <T> Provider<T> scope(final Key<T> key, final Provider<T> unscoped) {
    return new Provider<T>() {
      @SuppressWarnings("unchecked")
      public T get() {
        keys.put(key, state);
        Map<Class<?>, Object> map = getScopedObjectMap();
        
        if (!map.containsKey(state)) {
          T t = unscoped.get();
          map.put(state, t);
        }
        return (T) map.get(state);
      }
    };
  }
  
  private <T> Map<Class<?>, Object> getScopedObjectMap() {
    Map<Class<?>, Object> scopedObjects = values.get();
    if (scopedObjects == null) {
      scopedObjects = new HashMap<Class<?>, Object>();
      values.set(scopedObjects);
    }
    return scopedObjects;
  }
}