/*
 * Copyright 2020 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.
 */

//
// Lowers a module with loads and stores that access a 64-bit memory with
// one that works as-is on wasm32.
//
// TODO(wvo): make this run in parallel if needed.

#include "ir/bits.h"
#include "ir/import-utils.h"
#include "ir/localize.h"
#include "ir/utils.h"
#include "pass.h"
#include "wasm-builder.h"
#include "wasm.h"

namespace wasm {

static Name MEMORY_BASE("__memory_base");
static Name MEMORY_BASE32("__memory_base32");

static Name TABLE_BASE("__table_base");
static Name TABLE_BASE32("__table_base32");

static const Address k32GLimit(1ULL << 32);

struct Memory64Lowering : public WalkerPass<PostWalker<Memory64Lowering>> {

  bool refinalize = false;

  void wrapAddress64(Expression*& ptr,
                     Name memoryOrTableName,
                     bool isTable = false) {
    if (ptr->type == Type::unreachable) {
      return;
    }
    auto& module = *getModule();
    bool is64 = false;
    if (isTable) {
      is64 = module.getTable(memoryOrTableName)->is64();
    } else {
      is64 = module.getMemory(memoryOrTableName)->is64();
    }
    if (is64) {
      assert(ptr->type == Type::i64);
      ptr = Builder(module).makeUnary(UnaryOp::WrapInt64, ptr);
    }
  }

  void extendAddress64(Expression*& ptr,
                       Name memoryOrTableName,
                       bool isTable = false) {
    if (ptr->type == Type::unreachable) {
      return;
    }
    auto& module = *getModule();
    bool is64 = false;
    if (isTable) {
      is64 = module.getTable(memoryOrTableName)->is64();
    } else {
      is64 = module.getMemory(memoryOrTableName)->is64();
    }
    if (is64) {
      assert(ptr->type == Type::i64);
      ptr->type = Type::i32;
      ptr = Builder(module).makeUnary(UnaryOp::ExtendUInt32, ptr);
    }
  }

  void wrapTableAddress64(Expression*& ptr, Name tableName) {
    return wrapAddress64(ptr, tableName, true);
  }

  void extendTableAddress64(Expression*& ptr, Name tableName) {
    return extendAddress64(ptr, tableName, true);
  }

  template<typename T> void visitMemoryAccess(T* curr) {
    if (curr->offset < k32GLimit) {
      return wrapAddress64(curr->ptr, curr->memory);
    }
    Block* b =
      ChildLocalizer(curr, getFunction(), *getModule(), getPassOptions())
        .getChildrenReplacement();
    b->list.push_back(Builder(*getModule()).makeUnreachable());
    b->type = Type::unreachable;
    replaceCurrent(b);
    refinalize = true;
  }

  void visitLoad(Load* curr) { visitMemoryAccess(curr); }

  void visitStore(Store* curr) { visitMemoryAccess(curr); }

  void visitSIMDLoad(SIMDLoad* curr) { visitMemoryAccess(curr); }

  void visitSIMDLoadStoreLane(SIMDLoadStoreLane* curr) {
    visitMemoryAccess(curr);
  }

  void visitMemorySize(MemorySize* curr) {
    auto& module = *getModule();
    auto* memory = module.getMemory(curr->memory);
    if (memory->is64()) {
      auto* size = static_cast<Expression*>(curr);
      extendAddress64(size, curr->memory);
      curr->type = Type::i32;
      replaceCurrent(size);
    }
  }

  void visitMemoryGrow(MemoryGrow* curr) {
    auto& module = *getModule();
    auto* memory = module.getMemory(curr->memory);
    if (memory->is64()) {
      wrapAddress64(curr->delta, curr->memory);
      auto* size = static_cast<Expression*>(curr);
      // MemoryGrow returns -1 in case of failure.  We cannot just use
      // extend_32_u in this case so we handle it as follows:
      //
      // (if (result i64)
      //  (i32.eq (i32.const -1) (local.tee $tmp (memory.grow X)))
      //  (then
      //   (i64.const -1)
      //  )
      //  (else
      //   (i32.extend_32_u (local.get $tmp))
      //  )
      // )
      Builder builder(module);
      auto tmp = builder.addVar(getFunction(), Type::i32);
      Expression* isMinusOne =
        builder.makeBinary(EqInt32,
                           builder.makeConst(int32_t(-1)),
                           builder.makeLocalTee(tmp, size, Type::i32));
      auto* newSize = builder.makeLocalGet(tmp, Type::i32);
      builder.makeUnary(UnaryOp::ExtendUInt32, newSize);
      Expression* ifExp =
        builder.makeIf(isMinusOne,
                       builder.makeConst(int64_t(-1)),
                       builder.makeUnary(UnaryOp::ExtendUInt32, newSize));
      curr->type = Type::i32;
      replaceCurrent(ifExp);
    }
  }

  void visitMemoryInit(MemoryInit* curr) {
    wrapAddress64(curr->dest, curr->memory);
  }

  void visitMemoryFill(MemoryFill* curr) {
    wrapAddress64(curr->dest, curr->memory);
    wrapAddress64(curr->size, curr->memory);
  }

  void visitMemoryCopy(MemoryCopy* curr) {
    wrapAddress64(curr->dest, curr->destMemory);
    wrapAddress64(curr->source, curr->sourceMemory);
    wrapAddress64(curr->size, curr->destMemory);
  }

  void visitAtomicRMW(AtomicRMW* curr) { visitMemoryAccess(curr); }

  void visitAtomicCmpxchg(AtomicCmpxchg* curr) { visitMemoryAccess(curr); }

  void visitAtomicWait(AtomicWait* curr) { visitMemoryAccess(curr); }

  void visitAtomicNotify(AtomicNotify* curr) { visitMemoryAccess(curr); }

  void visitFunction(Function* func) {
    if (refinalize) {
      ReFinalize().walkFunctionInModule(func, getModule());
      refinalize = false;
    }
  }

  void visitDataSegment(DataSegment* segment) {
    auto& module = *getModule();

    // passive segments don't have any offset to adjust
    if (segment->isPassive || !module.getMemory(segment->memory)->is64()) {
      return;
    }

    if (auto* c = segment->offset->dynCast<Const>()) {
      c->value = Literal(static_cast<uint32_t>(c->value.geti64()));
      c->type = Type::i32;
    } else if (auto* get = segment->offset->dynCast<GlobalGet>()) {
      auto* g = module.getGlobal(get->name);
      if (g->imported() && g->base == MEMORY_BASE) {
        ImportInfo info(module);
        auto* memoryBase32 = info.getImportedGlobal(g->module, MEMORY_BASE32);
        if (!memoryBase32) {
          Builder builder(module);
          memoryBase32 = builder
                           .makeGlobal(MEMORY_BASE32,
                                       Type::i32,
                                       builder.makeConst(int32_t(0)),
                                       Builder::Immutable)
                           .release();
          memoryBase32->module = g->module;
          memoryBase32->base = MEMORY_BASE32;
          module.addGlobal(memoryBase32);
        }
        // Use this alternative import when initializing the segment.
        assert(memoryBase32);
        get->type = Type::i32;
        get->name = memoryBase32->name;
      }
    } else {
      WASM_UNREACHABLE("unexpected elem offset");
    }
  }

  void visitTableSize(TableSize* curr) {
    auto& module = *getModule();
    auto* table = module.getTable(curr->table);
    if (table->is64()) {
      auto* size = static_cast<Expression*>(curr);
      extendTableAddress64(size, curr->table);
      replaceCurrent(size);
    }
  }

  void visitTableGrow(TableGrow* curr) {
    auto& module = *getModule();
    auto* table = module.getTable(curr->table);
    if (table->is64()) {
      wrapTableAddress64(curr->delta, curr->table);
      auto* size = static_cast<Expression*>(curr);
      // TableGrow returns -1 in case of failure. We cannot just use
      // extend_32_u in this case so we handle it the same way as MemoryGrow:
      //
      // (if (result i64)
      //  (i32.eq (i32.const -1) (local.tee $tmp (table.grow X)))
      //  (then
      //   (i64.const -1)
      //  )
      //  (else
      //   (i64.extend_i32_u (local.get $tmp))
      //  )
      // )
      Builder builder(module);
      auto tmp = builder.addVar(getFunction(), Type::i32);
      Expression* isMinusOne =
        builder.makeBinary(EqInt32,
                           builder.makeConst(int32_t(-1)),
                           builder.makeLocalTee(tmp, size, Type::i32));
      auto* newSize = builder.makeLocalGet(tmp, Type::i32);
      Expression* ifExp =
        builder.makeIf(isMinusOne,
                       builder.makeConst(int64_t(-1)),
                       builder.makeUnary(UnaryOp::ExtendUInt32, newSize));
      curr->type = Type::i32;
      replaceCurrent(ifExp);
    }
  }

  void visitTableFill(TableFill* curr) {
    wrapTableAddress64(curr->dest, curr->table);
    wrapTableAddress64(curr->size, curr->table);
  }

  void visitTableCopy(TableCopy* curr) {
    wrapTableAddress64(curr->dest, curr->destTable);
    wrapTableAddress64(curr->source, curr->sourceTable);
    // The size type is i64 only when both tables are 64-bit.
    auto& module = *getModule();
    if (module.getTable(curr->destTable)->is64() &&
        module.getTable(curr->sourceTable)->is64()) {
      wrapAddress64(curr->size, curr->destTable, true);
    }
  }

  void visitTableInit(TableInit* curr) {
    wrapTableAddress64(curr->dest, curr->table);
  }

  void visitElemDrop(ElemDrop* curr) {}

  void visitCallIndirect(CallIndirect* curr) {
    wrapTableAddress64(curr->target, curr->table);
  }

  void visitElementSegment(ElementSegment* segment) {
    auto& module = *getModule();

    // Passive segments don't have any offset to update.
    if (segment->table.isNull() || !module.getTable(segment->table)->is64()) {
      return;
    }

    if (auto* c = segment->offset->dynCast<Const>()) {
      c->value = Literal(static_cast<uint32_t>(c->value.geti64()));
      c->type = Type::i32;
    } else if (auto* get = segment->offset->dynCast<GlobalGet>()) {
      auto* g = module.getGlobal(get->name);
      if (g->imported() && g->base == TABLE_BASE) {
        ImportInfo info(module);
        auto* memoryBase32 = info.getImportedGlobal(g->module, TABLE_BASE32);
        if (!memoryBase32) {
          Builder builder(module);
          memoryBase32 = builder
                           .makeGlobal(TABLE_BASE32,
                                       Type::i32,
                                       builder.makeConst(int32_t(0)),
                                       Builder::Immutable)
                           .release();
          memoryBase32->module = g->module;
          memoryBase32->base = TABLE_BASE32;
          module.addGlobal(memoryBase32);
        }
        // Use this alternative import when initializing the segment.
        assert(memoryBase32);
        get->type = Type::i32;
        get->name = memoryBase32->name;
      }
    } else {
      WASM_UNREACHABLE("unexpected elem offset");
    }
  }

  void run(Module* module) override {
    if (!module->features.has(FeatureSet::Memory64)) {
      return;
    }
    Super::run(module);
    // Don't modify the memories or tables themselves until after the traversal
    // since we that would require memories to be the last thing that get
    // visited, and we don't want to depend on that specific ordering.
    for (auto& memory : module->memories) {
      if (memory->is64()) {
        memory->addressType = Type::i32;
        if (memory->hasMax() && memory->max > memory->maxSize32()) {
          memory->max = memory->maxSize32();
        }
      }
    }
    for (auto& table : module->tables) {
      if (table->is64()) {
        table->addressType = Type::i32;
      }
    }
    module->features.disable(FeatureSet::Memory64);
  }
};

Pass* createMemory64LoweringPass() { return new Memory64Lowering(); }

} // namespace wasm