/*
 * Copyright 2026 WebAssembly Community Group participants
 *
 * Licensed 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.
 */

//
// Simple multi-allocation escape analysis, aiming to cleanly separate multiple
// allocations written to a single local.
//
// This pass identifies a simple case of refs not escaping, like Heap2Local, but
// allows more than a single allocation (struct.new), which is better, and
// only handles a far simpler-to-identify situation, which is worse. To see the
// motivation, consider overlapping allocations like this:
//
//    x = new Foo(42);
//    if (condition) x = new Foo(1337);
//    else x.value = -1;
//    bar(x.value, x.value);
//
// Heap2Local cannot optimize this due to the overlapping live ranges (it sees
// some x.value operations are not exclusive to any single allocation). But this
// pass can optimize into this:
//
//    x_value = (new Foo(42)).value;
//    if (condition) x_value = (new Foo(1337)).value;
//    else x_value = -1;
//    bar(x_value, x_value);
//
// Note how we now operate on x_value, a new local with data. Basically, each
// struct.get/struct.set is turned into an operation on that new local. This is
// safe despite the multiple allocations because the local x is always used
// immediately, as data: rather than store a reference, we can just focus on the
// data. That is, each struct.new just writes a bunch of fields, and it all
// could have been unboxed. After verifying that nothing observes x as a
// reference (no ref.eq), and it does not escape, we can optimize this.
// Heap2Local will then see 'untangled' allocations, each one immediately
// consumed, and remove them all.
//
// Note that Heap2Local could safely optimize in this example, even before our
// changes: while the `x.value` operations are not exclusive to allocations, the
// live ranges of the allocations join exactly there - they do not overlap in a
// dangerous way. But here is an example MultiHeap2Local can optimize that
// Heap2Local cannot:
//
//   (local.set $a
//     (struct.new $struct
//       (..value..)
//     )
//   )
//   (struct.get $ref 0
//     (select
//       (local.get $a)
//       (local.tee $a
//         (struct.new_default $struct)
//       )
//       (..condition..)
//     )
//   )
//
// The struct.get at the end may read ..value.. from set at the top, or it may
// read the default value from the tee at the bottom. But here we cannot
// optimize while using a single local `a_value` because of the live range
// overlap: both allocations are alive as the select executes, and it picks
// between them, so we must store both of them somewhere in the interim.
// Detecting such overlaps is not trivial as it involves locals as well as the
// wasm value stack, which is why this pass is useful, as it handles common and
// obviously-optimizable cases.
//

#include "ir/bits.h"
#include "ir/find_all.h"
#include "ir/local-graph.h"
#include "ir/properties.h"
#include "pass.h"
#include "wasm-builder.h"
#include "wasm-traversal.h"
#include "wasm.h"

namespace wasm {

namespace {

struct MultiHeap2Local : public WalkerPass<PostWalker<MultiHeap2Local>> {
  bool isFunctionParallel() override { return true; }

  std::unique_ptr<Pass> create() override {
    return std::make_unique<MultiHeap2Local>();
  }

  // Information about a local index, used to decide whether we can optimize.
  struct LocalInfo {
    // The number of gets for the local, tracked so that we can see if all the
    // local.gets are consumed by struct.get/sets.
    Index numGets = 0;

    // The number of sets for the local, tracked as we only optimize when we see
    // 2 or more local.sets (a single one is trivially optimized by Heap2Local,
    // so we would be doing pointless work).
    Index numSets = 0;

    // The number of tees for the local, tracked for the same reason as gets:
    // they flow out the reference, so we must see them consumed by struct
    // gets/sets.
    Index numTees = 0;

    // All struct.get/struct.set operations that read a get/tee of this local.
    std::vector<Expression*> structOperations;

    // Certain conditions prevent us from optimizing.
    bool unoptimizable = false;
  };

  // For each local index, its info.
  std::vector<LocalInfo> localInfos;

  // Checks if a local index is relevant for us: if it refers to a struct.
  bool isStructLocal(Index index) {
    auto type = getFunction()->getLocalType(index);
    return type.isRef() && type.getHeapType().isStruct();
  }

  void visitLocalGet(LocalGet* curr) {
    if (isStructLocal(curr->index)) {
      localInfos[curr->index].numGets++;
    }
  }

  void visitLocalSet(LocalSet* curr) {
    if (isStructLocal(curr->index)) {
      localInfos[curr->index].numSets++;
      if (curr->isTee()) {
        localInfos[curr->index].numTees++;
      }
      // All sets must write a struct.new, otherwise we cannot analyze this
      // local.
      if (!curr->value->is<StructNew>()) {
        localInfos[curr->index].unoptimizable = true;
      }
    }
  }

  void visitStructGet(StructGet* curr) { handleStructOperation(curr); }
  void visitStructSet(StructSet* curr) { handleStructOperation(curr); }

  template<typename T> void handleStructOperation(T* curr) {
    // Avoid unreachable code.
    if (curr->type == Type::unreachable) {
      return;
    }

    // See if the reference is one of the patterns we recognize.
    // TODO fallthrough, but careful of br_if which can escape the value
    if (auto* get = curr->ref->template dynCast<LocalGet>()) {
      localInfos[get->index].structOperations.push_back(curr);

      // Also handle the value for a struct.set: a struct.set has two inputs,
      // and if the value modifies the local, then two allocations can be
      // live at once:
      //
      //   (struct.set
      //     (local.get $x)
      //     (block
      //       (local.set $x ..)
      //       ..
      //
      // Here $x is put on the value stack, but then we write something new into
      // $x before the struct.set, which means the struct.set is writing to the
      // old allocation - one that no struct.get will later read. When we see
      // such live range overlaps, give up.
      if (auto* structSet = curr->template dynCast<StructSet>()) {
        for (auto* localSet : FindAll<LocalSet>(structSet->value).list) {
          if (localSet->index == get->index) {
            localInfos[get->index].unoptimizable = true;
            break;
          }
        }
      }
    } else if (auto* tee = curr->ref->template dynCast<LocalSet>()) {
      // TODO: Handle a tee for struct.set. This is a little complex to handle
      // below, and not really needed, since
      //
      //   (struct.set $A k
      //     (local.tee
      //       (struct.new $A ..)
      //     )
      //     (value)
      //   )
      //
      // can be folded into a single struct.new (that is, we are writing all the
      // fields, then writing an additional one, and other passes make us do
      // them all at once anyhow).
      if (curr->template is<StructSet>()) {
        localInfos[tee->index].unoptimizable = true;
      } else {
        // Otherwise, this is a valid struct operation.
        localInfos[tee->index].structOperations.push_back(curr);
      }
    }
    // TODO: handle more things flowing into struct.get, like an If
  }

  void doWalkFunction(Function* func) {
    localInfos.resize(func->getNumLocals());

    WalkerPass<PostWalker<MultiHeap2Local>>::doWalkFunction(func);

    pickLocalsToOptimize(func);

    if (!optimizedLocals.empty()) {
      optimize(func);
    }
  }

  // For each optimized local, maps each used field index to the new local
  // index that stores that field. Note that the inner map is sorted, both for
  // deterministic ordering when we create new locals, and also because the
  // number of used fields tends to be small (so hashing is overkill).
  std::unordered_map<Index, std::map<Index, Index>> optimizedLocals;

  void pickLocalsToOptimize(Function* func) {
    auto* module = getModule();
    Builder builder(*module);
    LazyLocalGraph localGraph(func, module);

    for (Index localIndex = 0; localIndex < localInfos.size(); ++localIndex) {
      auto& info = localInfos[localIndex];
      if (info.unoptimizable) {
        continue;
      }

      if (info.numSets < 2) {
        // We only care about the case of 2 or more sets for a local: otherwise,
        // Heap2Local can handle things perfectly for a single set.
        continue;
      }

      // Check if each local.get is consumed by a struct.get/set, which is the
      // case if the counts match: the sum of local.gets/local.tees (things
      // that flow out the reference) must match the number of struct.gets/sets
      // (things that consume the reference).
      if (info.numGets + info.numTees != info.structOperations.size()) {
        continue;
      }

      // Finally, if the initial value of the local from the function entry is
      // used, do not optimize. For a var, that means the default of zero, which
      // will trap when accessed, while for params it means an unknown value. We
      // therefore must always check for params, but for vars, if we know the
      // default cannot be accessed, we can skip this (which is the case for
      // a non-nullable local, and if trapsNeverHappen as a trap on the null
      // default can be ignored).
      auto localType = func->getLocalType(localIndex);
      auto canReadDefault =
        localType.isNullableRef() && !getPassOptions().trapsNeverHappen;
      if (func->isParam(localIndex) || canReadDefault) {
        bool ok = true;
        for (auto* op : info.structOperations) {
          Expression* ref;
          if (auto* get = op->dynCast<StructGet>()) {
            ref = get->ref;
          } else if (auto* set = op->dynCast<StructSet>()) {
            ref = set->ref;
          } else {
            WASM_UNREACHABLE("bad op");
          }
          if (auto* get = ref->dynCast<LocalGet>()) {
            if (localGraph.getSets(get).contains(nullptr)) {
              ok = false;
              break;
            }
          }
        }
        if (!ok) {
          continue;
        }
      }

      // We can optimize!

      // Gather all the struct field indexes operated on from this local.
      std::set<Index> usedFields;
      for (auto* op : info.structOperations) {
        usedFields.insert(getFieldIndex(op));
      }

      // Allocate new locals to store each field.
      auto& fields = localType.getHeapType().getStruct().fields;
      for (auto fieldIndex : usedFields) {
        optimizedLocals[localIndex][fieldIndex] =
          builder.addVar(func, fields[fieldIndex].type);
      }
    }
  }

  void optimize(Function* func) {
    struct Optimizer : PostWalker<Optimizer> {
      MultiHeap2Local& parent;

      Optimizer(MultiHeap2Local& parent) : parent(parent) {}

      void visitLocalSet(LocalSet* curr) {
        if (!curr->isTee() && parent.optimizedLocals.contains(curr->index)) {
          // local.set now stores the used fields (tee is handled below).
          replaceCurrent(makeFieldSets(curr));
        }
      }

      void visitStructGet(StructGet* curr) {
        if (auto* get = curr->ref->dynCast<LocalGet>()) {
          if (parent.optimizedLocals.contains(get->index)) {
            // struct.get of a local.get reads the field from a new local.
            replaceCurrent(
              getFieldFromLocal(get->index, curr->index, curr->signed_));
          }
        } else if (auto* tee = curr->ref->dynCast<LocalSet>()) {
          if (parent.optimizedLocals.contains(tee->index)) {
            // struct.get of a local.tee first writes the fields, then reads.
            replaceCurrent(makeFieldSets(tee, curr));
          }
        }
      }

      void visitStructSet(StructSet* curr) {
        if (auto* get = curr->ref->dynCast<LocalGet>()) {
          if (parent.optimizedLocals.contains(get->index)) {
            // struct.set of a local.get writes the field into a new local.
            replaceCurrent(
              setFieldToLocal(get->index, curr->index, curr->value));
          }
        }
      }

      Expression*
      getFieldFromLocal(Index localIndex, Index fieldIndex, bool signed_) {
        auto& field = getFunction()
                        ->getLocalType(localIndex)
                        .getHeapType()
                        .getStruct()
                        .fields[fieldIndex];
        // Read from the appropriate new local, handling sign extension as
        // needed.
        auto newLocal = parent.optimizedLocals[localIndex][fieldIndex];
        Builder builder(*getModule());
        Expression* ret =
          builder.makeLocalGet(newLocal, getFunction()->getLocalType(newLocal));
        return Bits::makePackedFieldGet(ret, field, signed_, *getModule());
      }

      Expression*
      setFieldToLocal(Index localIndex, Index fieldIndex, Expression* value) {
        // Write to the appropriate new local.
        auto newLocal = parent.optimizedLocals[localIndex][fieldIndex];
        Builder builder(*getModule());
        return builder.makeLocalSet(newLocal, value);
      }

      Expression* makeFieldSets(LocalSet* set, StructGet* teeParent = nullptr) {
        auto localIndex = set->index;
        auto localType = getFunction()->getLocalType(localIndex);
        auto& fields = localType.getHeapType().getStruct().fields;

        // Store the reference, then store struct.gets of all used fields in the
        // new locals.
        Builder builder(*getModule());
        std::vector<Expression*> list = {
          builder.makeLocalSet(localIndex, set->value)};
        for (auto [fieldIndex, newLocal] : parent.optimizedLocals[localIndex]) {
          auto& field = fields[fieldIndex];
          // Note that the sign does not matter here: we just store the bits,
          // and will handle the sign when reading back out (like in
          // Heap2Local).
          auto* sget =
            builder.makeStructGet(fieldIndex,
                                  builder.makeLocalGet(localIndex, localType),
                                  MemoryOrder::Unordered,
                                  field.type,
                                  false);
          list.push_back(builder.makeLocalSet(newLocal, sget));
        }

        // Finally, if this is a tee, then it was consumed by a struct.get, and
        // we can read that field and return it.
        if (teeParent) {
          list.push_back(getFieldFromLocal(
            localIndex, teeParent->index, teeParent->signed_));
        }
        return builder.makeBlock(list);
      }
    } optimizer(*this);
    optimizer.setModule(getModule());
    optimizer.walkFunction(func);
  }

  // Given a struct.get or struct.set, return the field index.
  Index getFieldIndex(Expression* curr) {
    if (auto* get = curr->dynCast<StructGet>()) {
      return get->index;
    } else if (auto* set = curr->dynCast<StructSet>()) {
      return set->index;
    } else {
      WASM_UNREACHABLE("bad op");
    }
  }
};

} // anonymous namespace

Pass* createMultiHeap2LocalPass() { return new MultiHeap2Local(); }

} // namespace wasm