package jetbrains.mps.kotlin.dataFlow;

/*Generated by MPS */

import org.jetbrains.annotations.ApiStatus;
import jetbrains.mps.kotlin.dataflow.instructions.DataflowEvaluator;
import java.util.HashMap;
import org.jetbrains.mps.openapi.model.SNode;
import jetbrains.mps.kotlin.smartcast.SmartCastType;
import org.jetbrains.annotations.Nullable;
import java.util.stream.Stream;
import jetbrains.mps.kotlin.dataflow.instructions.State;
import jetbrains.mps.baseLanguage.closures.runtime._FunctionTypes;
import jetbrains.mps.kotlin.dataflow.instructions.BooleanState;
import jetbrains.mps.lang.smodel.generator.smodelAdapter.SNodeOperations;
import jetbrains.mps.baseLanguage.closures.runtime.Wrappers;
import jetbrains.mps.kotlin.behavior.TypeReference;
import jetbrains.mps.kotlin.behavior.InferredTypeReference;
import jetbrains.mps.lang.dataFlow.framework.instructions.AbstractJumpInstruction;
import jetbrains.mps.kotlin.behavior.ConcreteTypeReference;
import org.jetbrains.mps.openapi.language.SConcept;
import jetbrains.mps.smodel.adapter.structure.MetaAdapterFactory;

@ApiStatus.Experimental
public class SmartCastEvaluator implements DataflowEvaluator<HashMap<SNode, SmartCastType>> {
  @Override
  public HashMap<SNode, SmartCastType> merge(HashMap<SNode, SmartCastType> left, @Nullable HashMap<SNode, SmartCastType> right) {
    if (right == null || right.isEmpty()) {
      return left;
    }

    final HashMap<SNode, SmartCastType> result = new HashMap<>(left);
    right.forEach((SNode key, final SmartCastType value) -> result.compute(key, (currentKey, currentValue) -> (currentValue == null ? value : value.union(currentValue))));

    return result;
  }

  private Stream<State<HashMap<SNode, SmartCastType>>> binaryOutcomes(final State<HashMap<SNode, SmartCastType>> previousState, final int instructionIndex, final _FunctionTypes._void_P2_E0<? super HashMap<SNode, SmartCastType>, ? super BooleanState> initializer) {
    // Map both true and false value
    return DataflowEvaluator.booleanStates().map((booleanState) -> {
      HashMap<SNode, SmartCastType> newStates = new HashMap<>(previousState.getData());
      initializer.invoke(newStates, booleanState);
      return State.of(newStates, instructionIndex + 1, booleanState);
    });
  }

  private Stream<State<HashMap<SNode, SmartCastType>>> singleOutcome(State<HashMap<SNode, SmartCastType>> previousState, int instruction, _FunctionTypes._void_P1_E0<? super HashMap<SNode, SmartCastType>> initializer) {
    // Map both true and false value
    HashMap<SNode, SmartCastType> newStates = new HashMap<>(previousState.getData());
    initializer.invoke(newStates);
    return Stream.of(State.of(newStates, instruction + 1));
  }

  @Override
  public Stream<State<HashMap<SNode, SmartCastType>>> acceptEquals(SNode leftNode, SNode rightNode, final boolean isNegation, int instructionIndex, State<HashMap<SNode, SmartCastType>> previousState) {
    final SNode left = FlowUtil.getLocal(leftNode);
    final SNode right = FlowUtil.getLocal(rightNode);

    if ((left != null) && (right != null)) {
      // local references equality
      final SmartCastType leftCast = previousState.getData().get(left);
      final SmartCastType rightCast = previousState.getData().get(right);
      if (leftCast != null && rightCast != null) {
        // Map both outcomes
        return binaryOutcomes(previousState, instructionIndex, (newState, booleanState) -> {
          boolean negate = isNegation ^ !(booleanState.booleanValue());
          newState.put(left, leftCast.intersect(negatedIfNeeded(rightCast, negate)));
          newState.put(right, rightCast.intersect(negatedIfNeeded(leftCast, negate)));
        });
      }
    } else if ((left == null) && (right == null)) {
      // no references
    } else if (SNodeOperations.isInstanceOf(rightNode, CONCEPTS.NullLiteral$jK) || SNodeOperations.isInstanceOf(leftNode, CONCEPTS.NullLiteral$jK)) {
      // one reference compared to null
      final SNode var = ((left != null) ? left : right);
      return binaryOutcomes(previousState, instructionIndex, (newState, booleanState) -> {
        final Wrappers._T<SmartCastType> addedType = new Wrappers._T<SmartCastType>();
        if (isNegation ^ !(booleanState.booleanValue())) {
          addedType.value = new SmartCastType(null, TypeReference.NULLABLE_NOTHING);
        } else {
          addedType.value = new SmartCastType(TypeReference.NULLABLE_NOTHING, null);
        }
        newState.compute(var, (_key, prevValue) -> addedType.value.intersect(prevValue));
      });
    }

    // default: pass
    return Stream.of(previousState);
  }

  private SmartCastType negatedIfNeeded(SmartCastType input, boolean negate) {
    if (negate) {
      return input.isNullable().swap();
    } else {
      return input;
    }
  }


  @Override
  public Stream<State<HashMap<SNode, SmartCastType>>> acceptAssign(SNode left, final SNode right, boolean isLeftAssignable, int instructionIndex, State<HashMap<SNode, SmartCastType>> previousState) {
    if (isLeftAssignable) {
      final SNode l = FlowUtil.getLocal(left);
      final SNode r = FlowUtil.getLocal(right);

      // TODO should break anterior lambda dataflow for this variable: have some object containing the current status for lambdas scope
      // TODO read and implement the smart cast sink concept https://kotlinlang.org/spec/type-inference.html#effectively-immutable-smart-cast-sinks
      if ((l != null) && (r != null)) {
        return singleOutcome(previousState, instructionIndex, (final HashMap<SNode, SmartCastType> newStates) -> {
          if (newStates.containsKey(r)) {
            newStates.compute(l, (_var, prevState) -> newStates.get(r).intersect(prevState));
          }
        });

      } else if ((l != null) && SNodeOperations.isInstanceOf(r, CONCEPTS.NullLiteral$jK)) {
        return singleOutcome(previousState, instructionIndex, (newStates) -> {
          final SmartCastType newState = new SmartCastType(TypeReference.NULLABLE_NOTHING, null);
          newStates.compute(l, (_key, prevValue) -> newState);
        });

      } else if ((l != null)) {
        return singleOutcome(previousState, instructionIndex, (newStates) -> {
          final SmartCastType newState = new SmartCastType(new InferredTypeReference(right), null);
          newStates.compute(l, (_key, prevValue) -> newState);
        });
      }
    }
    return Stream.of(previousState);
  }

  @Override
  public Stream<State<HashMap<SNode, SmartCastType>>> acceptBackedge(AbstractJumpInstruction instruction, State<HashMap<SNode, SmartCastType>> previousState) {
    return Stream.empty();
  }

  @Override
  public Stream<State<HashMap<SNode, SmartCastType>>> acceptIsType(SNode leftNode, final SNode right, final boolean negation, int instructionIndex, State<HashMap<SNode, SmartCastType>> previousState) {
    final SNode left = FlowUtil.getLocal(leftNode);
    if ((left != null)) {
      return binaryOutcomes(previousState, instructionIndex, (newStates, boolValue) -> {
        final Wrappers._T<SmartCastType> data = new Wrappers._T<SmartCastType>();
        boolean negate = negation ^ !(boolValue.booleanValue());
        ConcreteTypeReference reference = new ConcreteTypeReference(right);
        if (negate) {
          data.value = new SmartCastType(null, reference);
        } else {
          data.value = new SmartCastType(reference, null);
        }

        newStates.compute(left, (_key, state) -> data.value.intersect(state));
      });
    } else {
      return Stream.of(previousState);
    }
  }

  private static final class CONCEPTS {
    /*package*/ static final SConcept NullLiteral$jK = MetaAdapterFactory.getConcept(0x6b3888c1980244d8L, 0x8baff8e6c33ed689L, 0x28bef6d7551af3d5L, "jetbrains.mps.kotlin.structure.NullLiteral");
  }
}
