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.cassandra.harry.checker;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.function.BiConsumer;
  import java.util.function.BiPredicate;
  import java.util.function.Consumer;
  import java.util.function.Predicate;
  
  import org.apache.cassandra.harry.gen.EntropySource;
  import org.apache.cassandra.harry.gen.rng.JdkRandomEntropySource;
  
  public class ModelChecker<STATE, SUT>
  {
      private final List<StepExecutor<STATE, SUT>> steps;
      private final List<Precondition<STATE, SUT>> invariants;
      private Precondition<STATE, SUT> exitCondition;
      private Step<STATE, SUT> beforeAll;
      private Step<STATE, SUT> afterAll;
      private Pair<STATE, SUT> init;
  
      public ModelChecker()
      {
          steps = new ArrayList<>();
          invariants = new ArrayList<>();
      }
  
      public void run() throws Throwable
      {
          run(0, Long.MAX_VALUE, new JdkRandomEntropySource(System.currentTimeMillis()));
      }
  
      public void run(int minSteps, long maxSteps) throws Throwable
      {
          run(minSteps, maxSteps, new JdkRandomEntropySource(System.currentTimeMillis()));
      }
  
      public void run(int minSteps, long maxSteps, EntropySource rng) throws Throwable
      {
          assert init != null : "Initial condition is not specified";
  
          Ref<Pair<STATE, SUT>> state = new Ref<>(init, Pair.unchanged());
          if (beforeAll != null)
              state.map((s) -> beforeAll.next(s.l, s.r, rng));
  
          for (int i = 0; i < maxSteps; i++)
          {
              if (i > minSteps && exitCondition.test(state.get()))
                  return;
  
              // TODO: add randomisation / probability for triggering a specific step
              steps.get(rng.nextInt(steps.size())).execute(state, rng.derive());
              for (Precondition<STATE, SUT> invariant : invariants)
                  invariant.test(state.get());
          }
  
          if (afterAll != null)
              state.map((s) -> afterAll.next(s.l, s.r, rng));
      }
  
      public ModelChecker<STATE, SUT> init(STATE state, SUT sut)
      {
          this.init = new Pair<>(state, sut);
          return this;
      }
  
      public Simple init(STATE state)
      {
          Simple simple = new Simple();
          simple.init(state);
          return simple;
      }
  
      @SuppressWarnings("unused")
      public ModelChecker<STATE, SUT> beforeAll(Step<STATE, SUT> beforeAll)
      {
          this.beforeAll = beforeAll;
          return this;
      }
  
      @SuppressWarnings("unused")
      public ModelChecker<STATE, SUT> afterAll(Step<STATE, SUT> afterAll)
      {
          this.afterAll = afterAll;
          return this;
      }
  
      public ModelChecker<STATE, SUT> exitCondition(Precondition<STATE, SUT> precondition)
      {
          this.exitCondition = precondition;
          return this;
      }
  
      public ModelChecker<STATE, SUT> step(Precondition<STATE, SUT> precondition, Step<STATE, SUT> step)
      {
          steps.add((ref, rng) -> {
              ref.map(state -> {
                  if (!precondition.test(state))
                      return state;
  
                  return step.next(state.l, state.r, rng);
              });
          });
  
          return this;
      }
  
      public ModelChecker<STATE, SUT> invariant(Precondition<STATE, SUT> invariant)
      {
          invariants.add(invariant);
          return this;
      }
  
      public ModelChecker<STATE, SUT> step(Step<STATE, SUT> step)
      {
          return step(Precondition.alwaysTrue(), step);
      }
  
      public ModelChecker<STATE, SUT> step(StatePrecondition<STATE> precondition, ThrowingFunction<STATE, STATE> step)
      {
          steps.add((ref, rng) -> {
              ref.map(state -> {
                  if (!precondition.test(state.l))
                      return state;
  
                  STATE newState = step.apply(state.l);
                  if (state.l == newState)
                      return state;
  
                  return state.next(newState);
              });
          });
  
          return this;
      }
  
      public ModelChecker<STATE, SUT> step(StatePrecondition<STATE> precondition, ThrowingBiFunction<STATE, SUT, STATE> step)
      {
          steps.add((ref, rng) -> {
              ref.map(state -> {
                  if (!precondition.test(state.l))
                      return state;
  
                  STATE newState = step.apply(state.l, state.r);
                  if (state.l == newState)
                      return state;
  
                  return state.next(newState);
              });
          });
  
          return this;
      }
  
      public ModelChecker<STATE, SUT> step(ThrowingFunction<STATE, STATE> step)
      {
          return step((t, sut, rng) -> new Pair<>(step.apply(t), sut));
      }
  
      interface StepExecutor<STATE, SUT>
      {
          void execute(Ref<Pair<STATE, SUT>> state, EntropySource rng) throws Throwable;
      }
  
      public interface StatePrecondition<STATE>
      {
          boolean test(STATE state) throws Throwable;
      }
  
      public interface Precondition<STATE, SUT>
      {
          boolean test(STATE state, SUT sut) throws Throwable;
  
          default boolean test(Pair<STATE, SUT> state) throws Throwable
          {
              return test(state.l, state.r);
          }
  
          static <STATE, SUT> Precondition<STATE, SUT> alwaysTrue()
          {
              return (a, b) -> true;
          }
      }
  
      public interface Step<STATE, SUT>
      {
          Pair<STATE, SUT> next(STATE t, SUT sut, EntropySource rng) throws Throwable;
      }
  
      public interface ThrowingConsumer<I>
      {
          void accept(I t) throws Throwable;
      }
  
      public interface ThrowingBiConsumer<I1, I2>
      {
          void accept(I1 t1, I2 t2) throws Throwable;
      }
  
      public interface ThrowingFunction<I, O>
      {
          O apply(I t) throws Throwable;
      }
  
      public interface ThrowingBiFunction<I1, I2, O>
      {
          O apply(I1 t1, I2 t2) throws Throwable;
      }
  
      private static class Ref<T>
      {
          public T ref;
          private final T unchanged;
  
          public Ref(T init, T unchanged)
          {
              this.ref = init;
              this.unchanged = unchanged;
          }
  
          public T get()
          {
              return ref;
          }
  
          public void set(T v)
          {
              if (v == unchanged)
                  return;
              this.ref = v;
          }
  
          public void map(ThrowingFunction<T, T> fn) throws Throwable
          {
              set(fn.apply(ref));
          }
      }
  
      public static class Pair<L, R>
      {
          private static final Pair<?, ?> UNCHANGED = new Pair<>(null, null);
  
          public static <L, R> Pair<L, R> unchanged()
          {
              return (Pair<L, R>) UNCHANGED;
          }
  
          public final L l;
          public final R r;
  
          public Pair(L l, R r)
          {
              this.l = l;
              this.r = r;
          }
  
          public Pair<L, R> next(L state)
          {
              return new Pair<>(state, this.r);
          }
      }
  
      public class Simple
      {
          public Simple init(STATE state)
          {
              ModelChecker.this.init = new Pair<>(state, null);
              return this;
          }
  
          @SuppressWarnings("unused")
          public Simple beforeAll(ThrowingConsumer<STATE> beforeAll)
          {
              ModelChecker.this.beforeAll = (t, sut, rng) -> {
                  beforeAll.accept(t);
                  return Pair.unchanged();
              };
              return this;
          }
  
          public Simple beforeAll(ThrowingBiConsumer<STATE, EntropySource> beforeAll)
          {
              ModelChecker.this.beforeAll = (t, sut, rng) -> {
                  beforeAll.accept(t, rng);
                  return Pair.unchanged();
              };
              return this;
          }
  
          @SuppressWarnings("unused")
          public Simple beforeAll(ThrowingFunction<STATE, STATE> beforeAll)
          {
              ModelChecker.this.beforeAll = (t, sut, rng) -> new Pair<>(beforeAll.apply(t), sut);
              return this;
          }
  
          public Simple afterAll(ThrowingConsumer<STATE> afterAll)
          {
              ModelChecker.this.afterAll = (t, sut, rng) -> {
                  afterAll.accept(t);
                  return Pair.unchanged();
              };
              return this;
          }
  
          @SuppressWarnings("unused")
          public Simple afterAll(ThrowingFunction<STATE, STATE> afterAll)
          {
              ModelChecker.this.afterAll = (t, sut, rng) -> new Pair(afterAll.apply(t), sut);
              return this;
          }
  
          public Simple exitCondition(Predicate<STATE> precondition)
          {
              ModelChecker.this.exitCondition = (state, sut) -> precondition.test(state);
              return this;
          }
  
          @SuppressWarnings("unused")
          public Simple invariant(Predicate<STATE> invariant)
          {
              invariants.add((state, sut) -> invariant.test(state));
              return this;
          }
  
          public Simple step(ThrowingFunction<STATE, STATE> step)
          {
              ModelChecker.this.step((state, sut, rng) -> new Pair<>(step.apply(state), sut));
              return this;
          }
  
          public Simple step(ThrowingConsumer<STATE> step)
          {
              ModelChecker.this.step((state, sut, rng) -> {
                  step.accept(state);
                  return Pair.unchanged();
              });
              return this;
          }
  
          public Simple step(ThrowingBiConsumer<STATE, EntropySource> step)
          {
              ModelChecker.this.step((state, sut, rng) -> {
                  step.accept(state, rng);
                  return Pair.unchanged();
              });
              return this;
          }
  
          public Simple step(BiPredicate<STATE, EntropySource> precondition, BiConsumer<STATE, EntropySource> step)
          {
              ModelChecker.this.step((state, sut, rng) -> {
                  if (!precondition.test(state, rng))
                      return Pair.unchanged();
  
                  step.accept(state, rng);
                  return Pair.unchanged();
              });
  
              return this;
          }
  
          public Simple step(Predicate<STATE> precondition, Consumer<STATE> step)
          {
              ModelChecker.this.step((state, ignore) -> precondition.test(state),
                                     (t, sut, rng) -> {
                                         step.accept(t);
                                         return Pair.unchanged();
                                     });
  
              return this;
          }
  
          public void run() throws Throwable
          {
              ModelChecker.this.run();
          }
  
          public void run(EntropySource rng) throws Throwable
          {
              ModelChecker.this.run(0, Integer.MAX_VALUE, rng);
          }
  
          public void run(int minSteps, long maxSteps, EntropySource rng) throws Throwable
          {
              ModelChecker.this.run(minSteps, maxSteps, rng);
          }
      }
  
  }
  