/*
 * Licensed to Julian Hyde under one or more contributor license
 * agreements.  See the NOTICE file distributed with this work
 * for additional information regarding copyright ownership.
 * Julian Hyde 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 net.hydromatic.morel.compile;

import net.hydromatic.morel.ast.AstNode;
import net.hydromatic.morel.ast.Core;
import net.hydromatic.morel.type.TypeSystem;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Multimap;

import java.util.ArrayDeque;
import java.util.Collection;
import java.util.Deque;
import java.util.HashMap;
import java.util.Map;

/**
 * Shuttle that counts how many times each expression is used.
 */
public class Analyzer extends EnvVisitor {
  private final Map<Core.NamedPat, MutableUse> map;

  /** Creates an Analyzer. */
  private static Analyzer of(TypeSystem typeSystem, Environment env) {
    return new Analyzer(typeSystem, env, new HashMap<>(), new ArrayDeque<>());
  }

  /** Private constructor. */
  private Analyzer(TypeSystem typeSystem, Environment env,
      Map<Core.NamedPat, MutableUse> map, Deque<FromContext> fromStack) {
    super(typeSystem, env, fromStack);
    this.map = map;
  }

  /** Analyzes an expression. */
  public static Analysis analyze(TypeSystem typeSystem, Environment env,
      AstNode node) {
    final Analyzer analyzer = of(typeSystem, env);

    // Mark all top-level bindings so that they will not be removed
    if (node instanceof Core.NonRecValDecl) {
      analyzer.use(((Core.NonRecValDecl) node).pat).top = true;
    }
    node.accept(analyzer);
    return analyzer.result();
  }

  /** Returns the result of an analysis. */
  private Analysis result() {
    final ImmutableMap.Builder<Core.NamedPat, Use> b = ImmutableMap.builder();
    map.forEach((k, v) -> b.put(k, v.fix()));
    return new Analysis(b.build());
  }

  @Override protected Analyzer push(Environment env) {
    return new Analyzer(typeSystem, env, map, fromStack);
  }

  @Override protected void visit(Core.IdPat idPat) {
    use(idPat);
  }

  @Override public void visit(Core.Id id) {
    use(id.idPat).useCount++;
    super.visit(id);
  }

  /** Gets or creates a {@link MutableUse} for a given name. */
  private MutableUse use(Core.NamedPat name) {
    return map.computeIfAbsent(name, k -> new MutableUse());
  }

  @Override protected void visit(Core.NonRecValDecl valDecl) {
    super.visit(valDecl);
    if (isAtom(valDecl.exp)) {
      use(valDecl.pat).atomic = true;
    }
  }

  private static boolean isAtom(Core.Exp exp) {
    switch (exp.op) {
    case ID:
    case BOOL_LITERAL:
    case CHAR_LITERAL:
    case INT_LITERAL:
    case REAL_LITERAL:
    case STRING_LITERAL:
    case UNIT_LITERAL:
      return true;
    default:
      return false;
    }
  }

  @Override protected void visit(Core.Case case_) {
    case_.exp.accept(this);
    if (case_.matchList.size() == 1) {
      // When there is a single branch, we don't need to check for a single use
      // on multiple branches, so we can expedite.
      case_.matchList.get(0).accept(this);
    } else {
      // Create a multimap of all uses of bindings along the separate
      // branches. Example:
      //  case e of
      //    1 => a + c
      //  | 2 => a + b + a
      //  | _ => c
      //
      // a has use counts [1, 2] and is therefore MULTI_UNSAFE
      // b has use counts [1] and is therefore ONCE_SAFE
      // c has use counts [1, 1] and is therefore MULTI_SAFE
      final Multimap<Core.NamedPat, MutableUse> multimap =
          HashMultimap.create();
      final Map<Core.NamedPat, MutableUse> subMap = new HashMap<>();
      final Analyzer analyzer =
          new Analyzer(typeSystem, env, subMap, new ArrayDeque<>());
      case_.matchList.forEach(e -> {
        subMap.clear();
        e.accept(analyzer);
        subMap.forEach(multimap::put);
      });
      multimap.asMap().forEach((id, uses) -> {
        final MutableUse baseUse = use(id);
        int maxCount = MutableUse.max(uses);
        if (uses.size() > 1) {
          baseUse.parallel = true;
        }
        baseUse.useCount += maxCount;
      });
    }
  }

  /** How a binding (assignment of a value to a variable) is used. */
  public enum Use {
    /** Indicates that the binding cannot be inlined because recursively
     * refers to itself (or more precisely, is part of a recursive cycle
     * and has been chosen as the link to remove to break the cycle). */
    LOOP_BREAKER,

    /** Binding is not used. For a let (whether recursive or not), the binding
     * can be discarded. */
    DEAD,

    /** The binding occurs exactly once, and that occurrence is not inside a
     * lambda, nor is a constructor argument. Inlining is unconditionally safe;
     * it duplicates neither code nor work. */
    ONCE_SAFE,

    /** The binding is an atom (variable or literal). Regardless of how many
     * times it is used, inlining is unconditionally safe;
     * it duplicates neither code nor work. */
    ATOMIC,

    /** The binding occurs at most once in each of several distinct case
     * branches; none of these occurrences is inside a lambda. For example:
     *
     * <pre>{@code
     * case xs of
     *   [] => y + 1
     * | x :: xs => y + 2
     * }</pre>
     *
     * <p>In this expression, {@code y} occurs only once in each case branch.
     * Inlining {@code y} may duplicate code, but it will not duplicate work. */
    MULTI_SAFE,

    /** The binding occurs exactly once, but inside a lambda. Inlining will not
     * duplicate code, but it might duplicate work.
     *
     * <p>We must not inline an arbitrary expression inside a lambda, as the
     * following example (from GHC inlining section 2.2) shows:
     *
     * <pre>{@code
     * val f = fn x => E
     * val g = fn ys => map f ys
     * }</pre>
     *
     * <p>If we were to inline f inside g, thus:
     *
     * <pre>{@code
     * val g = fn ys => map (fn x => E) ys
     * }</pre>
     *
     * <p>no code is duplicated, but a small bounded amount of work is
     * duplicated, because the closure {@code fn x => E} must be allocated
     * each time {@code g} is called. */
    ONCE_UNSAFE,

    /** The binding may occur many times, including inside lambdas. */
    MULTI_UNSAFE,
  }

  /** Work space where the uses of a binding are counted. When all the uses
   * have been found, call {@link #fix} to convert this into a {@link Use}. */
  private static class MutableUse {
    boolean top;
    boolean atomic;
    boolean parallel;
    int useCount;

    static int max(Collection<MutableUse> uses) {
      int max = 0;
      for (MutableUse use : uses) {
        max = Math.max(max, use.useCount);
      }
      return max;
    }

    @Override public String toString() {
      return "[" + useCount + (parallel ? " parallel]" : "]");
    }

    Use fix() {
      return top ? Use.MULTI_UNSAFE
          : useCount == 0 ? Use.DEAD
          : atomic ? Use.ATOMIC
          : useCount == 1 ? (parallel ? Use.MULTI_SAFE : Use.ONCE_SAFE)
          : Use.MULTI_UNSAFE;
    }
  }

  /** Result of an analysis. */
  public static class Analysis {
    public final ImmutableMap<Core.NamedPat, Use> map;

    Analysis(ImmutableMap<Core.NamedPat, Use> map) {
      this.map = map;
    }
  }
}

// End Analyzer.java
