content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF 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 org.apache.calcite.util;
  
  import org.checkerframework.checker.nullness.qual.NonNull;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  import java.util.function.Supplier;
  
  import static org.apache.calcite.linq4j.Nullness.castNonNull;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Thread-local variable that returns a handle that can be closed.
   *
   * @param <T> Value type
   */
  public abstract class TryThreadLocal<T> extends ThreadLocal<@Nullable T> {
    /** Creates a TryThreadLocal with a fixed initial value.
     *
     * @param initialValue Initial value
     */
    public static <S> TryThreadLocal<S> of(S initialValue) {
      return new FixedTryThreadLocal<>(initialValue);
    }
  
    /** Creates a TryThreadLocal with a fixed initial value
     * whose values are never null.
     *
     * <p>The value returned from {@link #get} is never null;
     * the initial value must not be null;
     * you must not call {@link #set(Object)} with a null value.
     *
     * @param initialValue Initial value
     */
    public static <S> TryThreadLocal<@NonNull S> ofNonNull(S initialValue) {
      return new NonNullFixedTryThreadLocal<>(
          requireNonNull(initialValue, "initialValue"));
    }
  
    /** Creates a TryThreadLocal with a supplier for the initial value.
     *
     * <p>The value returned from {@link #get} is never null;
     * the supplier must never return null;
     * you must not call {@link #set(Object)} with a null value.
     *
     * @param supplier Supplier
     */
    public static <S> TryThreadLocal<@NonNull S> withInitial(
        Supplier<? extends @NonNull S> supplier) {
      return new SuppliedTryThreadLocal<>(supplier);
    }
  
    @Override public T get() {
      return castNonNull(super.get());
    }
  
    /** Assigns the value as {@code value} for the current thread.
     * Returns a {@link Memo} which, when closed, will assign the value
     * back to the previous value. */
    public Memo push(T value) {
      final T previous = get();
      set(value);
      return () -> restoreTo(previous);
    }
  
    /** Sets the value back to a previous value. */
    protected abstract void restoreTo(T previous);
  
    /** Performs an action with this ThreadLocal set to a particular value
     * in this thread, and restores the previous value afterward.
     *
     * <p>This method is named after the Standard ML {@code let} construct,
     * for example {@code let val x = 1 in x + 2 end}. */
    public void letIn(T t, Runnable runnable) {
      final T previous = get();
      if (previous == t) {
        runnable.run();
      } else {
        try {
          set(t);
          runnable.run();
        } finally {
          restoreTo(previous);
        }
      }
    }
  
    /** Calls a Supplier with this ThreadLocal set to a particular value,
     * in this thread, and restores the previous value afterward.
     *
     * <p>This method is named after the Standard ML {@code let} construct,
     * for example {@code let val x = 1 in x + 2 end}. */
    public <R> R letIn(T t, Supplier<R> supplier) {
      final T previous = get();
      if (previous == t) {
        return supplier.get();
      } else {
        try {
          set(t);
          return supplier.get();
        } finally {
          restoreTo(previous);
        }
      }
    }
  
    /** Remembers to set the value back. */
    public interface Memo extends AutoCloseable {
      /** Sets the value back; never throws. */
      @Override void close();
    }
  
    /** Implementation of {@link org.apache.calcite.util.TryThreadLocal}
     * with a fixed initial value.
     *
     * @param <T> Value type */
    private static class FixedTryThreadLocal<T> extends TryThreadLocal<T> {
      private final T initialValue;
  
      protected FixedTryThreadLocal(T initialValue) {
        this.initialValue = initialValue;
      }
  
      /**
       * {@inheritDoc}
       *
       * <p>It is important that this method is final.
       * This ensures that a subclass does not choose a different initial
       * value. Then the close logic can detect whether the previous value was
       * equal to the initial value.
       */
      @Override protected final T initialValue() {
        return initialValue;
      }
  
      /** Sets the value back to a previous value.
       *
       * <p>If the previous value was {@link #initialValue}, calls
       * {@link #remove()}. There's no way to tell whether {@link #set} has
       * been called previously, but the effect is the same. */
      @Override protected void restoreTo(T previous) {
        if (previous == initialValue) {
          remove();
        } else {
          set(previous);
        }
      }
    }
  
    /** Implementation of {@link org.apache.calcite.util.TryThreadLocal}
     * with a fixed initial value.
     *
     * @param <T> Value type */
    private static class NonNullFixedTryThreadLocal<T>
        extends FixedTryThreadLocal<T> {
      private NonNullFixedTryThreadLocal(T initialValue) {
        super(requireNonNull(initialValue, "initialValue"));
      }
  
      @Override public void set(@Nullable T value) {
        super.set(requireNonNull(value, "value"));
      }
  
      @Override public T get() {
        return requireNonNull(super.get());
      }
    }
  
    /** Implementation of {@link org.apache.calcite.util.TryThreadLocal}
     * whose initial value comes from a supplier.
     *
     * @param <T> Value type */
    private static class SuppliedTryThreadLocal<T> extends TryThreadLocal<T> {
      private final Supplier<? extends @NonNull T> supplier;
  
      SuppliedTryThreadLocal(Supplier<? extends @NonNull T> supplier) {
        this.supplier = requireNonNull(supplier, "supplier");
      }
  
      @Override protected @NonNull T initialValue() {
        return requireNonNull(supplier.get(), "supplier returned null");
      }
  
      @Override public void set(@Nullable T value) {
        super.set(requireNonNull(value, "value"));
      }
  
      @Override protected void restoreTo(T previous) {
        // If the thread had no value before they called 'push', should we call
        // 'remove()' here? No, for two reasons.
        //
        // First, it's not possible to know whether there was a value.
        // (ThreadLocal.isPresent() is package-protected.)
        //
        // Second, it may be what the user wants. If each 'restoreTo' call invokes
        // 'remove', then the next call to 'push' will invoke the supplier again.
        // Sometimes the user doesn't want to pay the initialization cost multiple
        // times, or to lose the state in the initialized object.
        set(previous);
      }
    }
  }
  