com.simiacryptus.mindseye.lang.DoubleBufferSet Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2019 by Andrew Charneski.
 *
 * The author 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.
 */

package com.simiacryptus.mindseye.lang;

import com.simiacryptus.lang.ref.ReferenceCounting;
import com.simiacryptus.lang.ref.ReferenceCountingBase;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public abstract class DoubleBufferSet> extends ReferenceCountingBase {
  static final Logger log = LoggerFactory.getLogger(DoubleBufferSet.class);

  @Nonnull
  protected final HashMap map = new HashMap<>();

  public DoubleBufferSet() {
  }

  public DoubleBufferSet(@Nonnull final DoubleBufferSet toCopy) {
    this(toCopy.map);
  }

  public DoubleBufferSet(@Nonnull final Map collect) {
    collect.forEach((k, v) -> {
      assert null != k;
      assert null != v;
    });
    synchronized (collect) {
      map.putAll(collect);
      map.forEach((k, v) -> {
        if (k instanceof ReferenceCounting) ((ReferenceCounting) k).addRef();
        v.addRef();
      });
    }
  }

  @Nonnull
  @SuppressWarnings("unchecked")
  public DoubleBufferSet copy() {
    return map(x -> (T) x.copy());
  }

  protected abstract T factory(final K layer, final double[] target);

  public T get(final K layer, final double[] ptr) {
    final T delta = get(layer, () -> factory(layer, ptr));
    assert delta.key.equals(layer);
    assert delta.target == ptr;
    return delta;
  }

  private T get(@Nullable final K layer, @Nullable final Supplier factory) {
    if (null == map) throw new IllegalArgumentException();
    if (null == factory) throw new IllegalArgumentException();
    if (null == layer) throw new IllegalArgumentException();
    synchronized (map) {
      T v = map.computeIfAbsent(layer, l -> {
        if (l instanceof ReferenceCounting) ((ReferenceCounting) l).addRef();
        T delta = factory.get();
        assert null != delta;
        if (log.isDebugEnabled())
          log.debug(String.format("Init key buffer for %s - %s params", l.getClass(), delta.target.length));
        return delta;
      });
      v.addRef();
      return v;
    }
  }

  public T get(final K layer, @Nonnull final Tensor ptr) {
    return get(layer, ptr.getData());
  }

  @Nonnull
  public Map getMap() {
    return Collections.unmodifiableMap(map);
  }

  @Nonnull
  public DoubleBufferSet map(@Nonnull final Function mapper) {
    @Nonnull final DoubleBufferSet parent = this;
    Stream> stream = map.entrySet().stream();
    if (map.size() > 100) {
      stream = stream.parallel();
    }
    final Map newMap = stream.collect(Collectors.toMap(e -> e.getKey(), e -> mapper.apply(e.getValue())));
    @Nonnull Delegate delegate = new Delegate(parent, newMap);
    newMap.values().forEach(x -> x.freeRef());
    return delegate;
  }

  public Stream stream() {
    return map.values().stream().filter(n -> null != n).distinct().sorted(Comparator.comparing(y -> System.identityHashCode(y.target)));
  }

  @Override
  protected void _free() {
    map.forEach((k, v) -> {
      if (k instanceof ReferenceCounting) ((ReferenceCounting) k).freeRef();
      v.freeRef();
    });
//    map.clear();
  }

  protected class Delegate extends DoubleBufferSet {
    private final DoubleBufferSet parent;

    public Delegate(final DoubleBufferSet parent) {
      this(parent, new HashMap<>());
    }

    public Delegate(final DoubleBufferSet parent, @Nonnull final Map newMap) {
      super(newMap);
      this.parent = parent;
    }

    @Override
    protected T factory(final K layer, final double[] target) {
      return parent.factory(layer, target);
    }
  }
}