// Copyright 2005-2024 Google LLC
//
// 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.
//
// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.
//
// NgramFst implements a n-gram language model based upon the LOUDS data
// structure. Please refer to "Unary Data Structures for Language Models"
// http://research.google.com/pubs/archive/37218.pdf

#ifndef FST_EXTENSIONS_NGRAM_NGRAM_FST_H_
#define FST_EXTENSIONS_NGRAM_NGRAM_FST_H_

#include <sys/types.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <ios>
#include <iostream>
#include <istream>
#include <memory>
#include <ostream>
#include <queue>
#include <string>
#include <utility>
#include <vector>

#include <fst/compat.h>
#include <fst/log.h>
#include <fst/extensions/ngram/bitmap-index.h>
#include <fst/arcsort.h>
#include <fst/expanded-fst.h>
#include <fstream>
#include <fst/fst.h>
#include <fst/fstlib.h>
#include <fst/mapped-file.h>
#include <fst/matcher.h>
#include <fst/properties.h>
#include <fst/util.h>
#include <fst/vector-fst.h>

namespace fst {
template <class A>
class NGramFst;
template <class A>
class NGramFstMatcher;

// Instance data containing mutable state for bookkeeping repeated access to
// the same state.
template <class A>
struct NGramFstInst {
  typedef typename A::Label Label;
  typedef typename A::StateId StateId;
  typedef typename A::Weight Weight;
  StateId state_;
  size_t num_futures_;
  size_t offset_;
  size_t node_;
  StateId node_state_;
  std::vector<Label> context_;
  StateId context_state_;
  NGramFstInst()
      : state_(kNoStateId),
        node_state_(kNoStateId),
        context_state_(kNoStateId) {}
};

namespace internal {

// Implementation class for LOUDS based NgramFst interface.
template <class A>
class NGramFstImpl : public FstImpl<A> {
  using FstImpl<A>::SetInputSymbols;
  using FstImpl<A>::SetOutputSymbols;
  using FstImpl<A>::SetType;
  using FstImpl<A>::WriteHeader;

  friend class ArcIterator<NGramFst<A>>;
  friend class NGramFstMatcher<A>;

 public:
  using FstImpl<A>::InputSymbols;
  using FstImpl<A>::SetProperties;
  using FstImpl<A>::Properties;

  typedef A Arc;
  typedef typename A::Label Label;
  typedef typename A::StateId StateId;
  typedef typename A::Weight Weight;

  NGramFstImpl() {
    SetType("ngram");
    SetInputSymbols(nullptr);
    SetOutputSymbols(nullptr);
    SetProperties(kStaticProperties);
  }

  NGramFstImpl(const Fst<A> &fst, std::vector<StateId> *order_out);

  explicit NGramFstImpl(const Fst<A> &fst) : NGramFstImpl(fst, nullptr) {}

  NGramFstImpl(const NGramFstImpl &other) {
    FSTERROR() << "Copying NGramFst Impls is not supported, use safe = false.";
    SetProperties(kError, kError);
  }

  static NGramFstImpl<A> *Read(std::istream &strm, const FstReadOptions &opts) {
    auto impl = std::make_unique<NGramFstImpl<A>>();
    FstHeader hdr;
    if (!impl->ReadHeader(strm, opts, kMinFileVersion, &hdr)) return nullptr;
    uint64_t num_states, num_futures, num_final;
    const size_t offset =
        sizeof(num_states) + sizeof(num_futures) + sizeof(num_final);
    // Peek at num_states and num_futures to see how much more needs to be read.
    strm.read(reinterpret_cast<char *>(&num_states), sizeof(num_states));
    strm.read(reinterpret_cast<char *>(&num_futures), sizeof(num_futures));
    strm.read(reinterpret_cast<char *>(&num_final), sizeof(num_final));
    size_t size = Storage(num_states, num_futures, num_final);
    std::unique_ptr<MappedFile> data_region(MappedFile::Allocate(size));
    char *data = static_cast<char *>(data_region->mutable_data());
    // Copy num_states, num_futures and num_final back into data.
    memcpy(data, reinterpret_cast<char *>(&num_states), sizeof(num_states));
    memcpy(data + sizeof(num_states), reinterpret_cast<char *>(&num_futures),
           sizeof(num_futures));
    memcpy(data + sizeof(num_states) + sizeof(num_futures),
           reinterpret_cast<char *>(&num_final), sizeof(num_final));
    strm.read(data + offset, size - offset);
    if (strm.fail()) return nullptr;
    impl->Init(data, std::move(data_region));
    return impl.release();
  }

  bool Write(std::ostream &strm, const FstWriteOptions &opts) const {
    FstHeader hdr;
    hdr.SetStart(Start());
    hdr.SetNumStates(num_states_);
    WriteHeader(strm, opts, kFileVersion, &hdr);
    strm.write(data_, StorageSize());
    return !strm.fail();
  }

  StateId Start() const { return start_; }

  Weight Final(StateId state) const {
    if (final_index_.Get(state)) {
      return final_probs_[final_index_.Rank1(state)];
    } else {
      return Weight::Zero();
    }
  }

  size_t NumArcs(StateId state, NGramFstInst<A> *inst = nullptr) const {
    if (inst == nullptr) {
      const std::pair<size_t, size_t> zeros =
          (state == 0) ? select_root_ : future_index_.Select0s(state);
      return zeros.second - zeros.first - 1;
    }
    SetInstFuture(state, inst);
    return inst->num_futures_ + ((state == 0) ? 0 : 1);
  }

  size_t NumInputEpsilons(StateId state) const {
    // State 0 has no parent, thus no backoff.
    if (state == 0) return 0;
    return 1;
  }

  size_t NumOutputEpsilons(StateId state) const {
    return NumInputEpsilons(state);
  }

  StateId NumStates() const { return num_states_; }

  void InitStateIterator(StateIteratorData<A> *data) const {
    data->base = nullptr;
    data->nstates = num_states_;
  }

  static size_t Storage(uint64_t num_states, uint64_t num_futures,
                        uint64_t num_final) {
    uint64_t b64;
    Weight weight;
    Label label;
    size_t offset =
        sizeof(num_states) + sizeof(num_futures) + sizeof(num_final);
    offset +=
        sizeof(b64) * (BitmapIndex::StorageSize(num_states * 2 + 1) +
                       BitmapIndex::StorageSize(num_futures + num_states + 1) +
                       BitmapIndex::StorageSize(num_states));
    offset += (num_states + 1) * sizeof(label) + num_futures * sizeof(label);
    // Pad for alignemnt, see
    // http://en.wikipedia.org/wiki/Data_structure_alignment#Computing_padding
    offset = (offset + sizeof(weight) - 1) & ~(sizeof(weight) - 1);
    offset += (num_states + 1) * sizeof(weight) + num_final * sizeof(weight) +
              (num_futures + 1) * sizeof(weight);
    return offset;
  }

  void SetInstFuture(StateId state, NGramFstInst<A> *inst) const {
    if (inst->state_ != state) {
      inst->state_ = state;
      const std::pair<size_t, size_t> zeros = future_index_.Select0s(state);
      inst->num_futures_ = zeros.second - zeros.first - 1;
      inst->offset_ = future_index_.Rank1(zeros.first + 1);
    }
  }

  void SetInstNode(NGramFstInst<A> *inst) const {
    if (inst->node_state_ != inst->state_) {
      inst->node_state_ = inst->state_;
      inst->node_ = context_index_.Select1(inst->state_);
    }
  }

  void SetInstContext(NGramFstInst<A> *inst) const {
    SetInstNode(inst);
    if (inst->context_state_ != inst->state_) {
      inst->context_state_ = inst->state_;
      inst->context_.clear();
      size_t node = inst->node_;
      while (node != 0) {
        const size_t rank1 = context_index_.Rank1(node);
        const size_t rank0 = node - rank1;
        inst->context_.push_back(context_words_[rank1]);
        node = context_index_.Select1(rank0 - 1);
      }
    }
  }

  // Access to the underlying representation
  const char *GetData(size_t *data_size) const {
    *data_size = StorageSize();
    return data_;
  }

  void Init(const char *data,
            std::unique_ptr<MappedFile> data_region);

  const std::vector<Label> &GetContext(StateId s, NGramFstInst<A> *inst) const {
    SetInstFuture(s, inst);
    SetInstContext(inst);
    return inst->context_;
  }

  size_t StorageSize() const {
    return Storage(num_states_, num_futures_, num_final_);
  }

  void GetStates(const std::vector<Label> &context,
                 std::vector<StateId> *states) const;

 private:
  StateId Transition(const std::vector<Label> &context, Label future) const;

  // Properties always true for this Fst class.
  static constexpr uint64_t kStaticProperties =
      kAcceptor | kIDeterministic | kODeterministic | kEpsilons | kIEpsilons |
      kOEpsilons | kILabelSorted | kOLabelSorted | kWeighted | kCyclic |
      kInitialAcyclic | kNotTopSorted | kAccessible | kCoAccessible |
      kNotString | kExpanded;
  // Current file format version.
  static constexpr int kFileVersion = 4;
  // Minimum file format version supported.
  static constexpr int kMinFileVersion = 4;

  std::unique_ptr<MappedFile> data_region_;
  const char *data_ = nullptr;  // Not owned.
  StateId start_ = fst::kNoStateId;
  uint64_t num_states_ = 0;
  uint64_t num_futures_ = 0;
  uint64_t num_final_ = 0;
  std::pair<size_t, size_t> select_root_;
  const Label *root_children_ = nullptr;
  // borrowed references
  const uint64_t *context_ = nullptr;
  const uint64_t *future_ = nullptr;
  const uint64_t *final_ = nullptr;
  const Label *context_words_ = nullptr;
  const Label *future_words_ = nullptr;
  const Weight *backoff_ = nullptr;
  const Weight *final_probs_ = nullptr;
  const Weight *future_probs_ = nullptr;
  // Uses all operations.
  BitmapIndex context_index_;
  // Uses Select0 and Rank1.
  BitmapIndex future_index_;
  // Uses Get and Rank1. This wastes space if there are no or few final
  // states, but it's also small. TODO(jrosenstock): Look at EliasFanoArray.
  BitmapIndex final_index_;
};

template <typename A>
inline void NGramFstImpl<A>::GetStates(
    const std::vector<Label> &context,
    std::vector<typename A::StateId> *states) const {
  states->clear();
  states->push_back(0);
  typename std::vector<Label>::const_reverse_iterator cit = context.rbegin();
  const Label *children = root_children_;
  size_t num_children = select_root_.second - 2;
  const Label *loc = std::lower_bound(children, children + num_children, *cit);
  if (loc == children + num_children || *loc != *cit) return;
  size_t node = 2 + loc - children;
  size_t node_rank = context_index_.Rank1(node);
  states->push_back(node_rank);
  if (context.size() == 1) return;
  std::pair<size_t, size_t> zeros =
      node_rank == 0 ? select_root_ : context_index_.Select0s(node_rank);
  size_t first_child = zeros.first + 1;
  ++cit;
  if (context_index_.Get(first_child) != false) {
    size_t last_child = zeros.second - 1;
    while (cit != context.rend()) {
      children = context_words_ + context_index_.Rank1(first_child);
      loc = std::lower_bound(children, children + last_child - first_child + 1,
                             *cit);
      if (loc == children + last_child - first_child + 1 || *loc != *cit) {
        break;
      }
      ++cit;
      node = first_child + loc - children;
      node_rank = context_index_.Rank1(node);
      states->push_back(node_rank);
      zeros =
          node_rank == 0 ? select_root_ : context_index_.Select0s(node_rank);
      first_child = zeros.first + 1;
      if (context_index_.Get(first_child) == false) break;
      last_child = zeros.second - 1;
    }
  }
}

}  // namespace internal

/*****************************************************************************/
template <class A>
class NGramFst : public ImplToExpandedFst<internal::NGramFstImpl<A>> {
  using Base = ImplToExpandedFst<internal::NGramFstImpl<A>>;
  friend class ArcIterator<NGramFst<A>>;
  friend class NGramFstMatcher<A>;

 public:
  typedef A Arc;
  typedef typename A::StateId StateId;
  typedef typename A::Label Label;
  typedef typename A::Weight Weight;
  using typename Base::Impl;

  explicit NGramFst(const Fst<A> &dst)
      : Base(std::make_shared<Impl>(dst, nullptr)) {}

  NGramFst(const Fst<A> &fst, std::vector<StateId> *order_out)
      : Base(std::make_shared<Impl>(fst, order_out)) {}

  // Because the NGramFstImpl is a const stateless data structure, there
  // is never a need to do anything beside copy the reference.
  NGramFst(const NGramFst<A> &fst, bool safe = false) : Base(fst, false) {}

  NGramFst() : Base(std::make_shared<Impl>()) {}

  // Non-standard constructor to initialize NGramFst directly from data. Caller
  // maintains ownership of data, which must outlive the NGramFst.
  explicit NGramFst(const char *data) : Base(std::make_shared<Impl>()) {
    GetMutableImpl()->Init(data, /*data_region=*/nullptr);
  }

  // Get method that gets the data associated with Init().
  const char *GetData(size_t *data_size) const {
    return GetImpl()->GetData(data_size);
  }

  const std::vector<Label> GetContext(StateId s) const {
    return GetImpl()->GetContext(s, &inst_);
  }

  // Consumes as much as possible of context from right to left, returns the
  // the states corresponding to the increasingly conditioned input sequence.
  void GetStates(const std::vector<Label> &context,
                 std::vector<StateId> *state) const {
    return GetImpl()->GetStates(context, state);
  }

  size_t NumArcs(StateId s) const override {
    return GetImpl()->NumArcs(s, &inst_);
  }

  NGramFst<A> *Copy(bool safe = false) const override {
    return new NGramFst(*this, safe);
  }

  static NGramFst<A> *Read(std::istream &strm, const FstReadOptions &opts) {
    Impl *impl = Impl::Read(strm, opts);
    return impl ? new NGramFst<A>(std::shared_ptr<Impl>(impl)) : nullptr;
  }

  static NGramFst<A> *Read(const std::string &source) {
    if (!source.empty()) {
      std::ifstream strm(source,
                              std::ios_base::in | std::ios_base::binary);
      if (!strm.good()) {
        LOG(ERROR) << "NGramFst::Read: Can't open file: " << source;
        return nullptr;
      }
      return Read(strm, FstReadOptions(source));
    } else {
      return Read(std::cin, FstReadOptions("standard input"));
    }
  }

  bool Write(std::ostream &strm, const FstWriteOptions &opts) const override {
    return GetImpl()->Write(strm, opts);
  }

  bool Write(const std::string &source) const override {
    return Fst<A>::WriteFile(source);
  }

  inline void InitStateIterator(StateIteratorData<A> *data) const override {
    GetImpl()->InitStateIterator(data);
  }

  inline void InitArcIterator(StateId s,
                              ArcIteratorData<A> *data) const override;

  MatcherBase<A> *InitMatcher(MatchType match_type) const override {
    return new NGramFstMatcher<A>(this, match_type);
  }

  size_t StorageSize() const { return GetImpl()->StorageSize(); }

  static bool HasRequiredProps(const Fst<A> &fst) {
    static const auto props =
        kAcceptor | kIDeterministic | kILabelSorted | kIEpsilons | kAccessible;
    return fst.Properties(props, true) == props;
  }

  static bool HasRequiredStructure(const Fst<A> &fst) {
    if (!HasRequiredProps(fst)) {
      return false;
    }
    typename A::StateId unigram = fst.Start();
    while (true) {  // Follows epsilon arc chain to find unigram state.
      if (unigram == fst::kNoStateId) return false;  // No unigram state.
      typename fst::ArcIterator<Fst<A>> aiter(fst, unigram);
      if (aiter.Done() || aiter.Value().ilabel != 0) break;
      unigram = aiter.Value().nextstate;
      aiter.Next();
    }
    // Other requirement: all states other than unigram an epsilon arc.
    for (fst::StateIterator<Fst<A>> siter(fst); !siter.Done();
         siter.Next()) {
      const typename A::StateId &state = siter.Value();
      fst::ArcIterator<Fst<A>> aiter(fst, state);
      if (state != unigram) {
        if (aiter.Done()) return false;
        if (aiter.Value().ilabel != 0) return false;
        aiter.Next();
        if (!aiter.Done() && aiter.Value().ilabel == 0) return false;
      }
    }
    return true;
  }

 private:
  using Base::GetImpl;
  using Base::GetMutableImpl;

  explicit NGramFst(std::shared_ptr<Impl> impl) : Base(impl) {}

  mutable NGramFstInst<A> inst_;
};

template <class A>
inline void NGramFst<A>::InitArcIterator(StateId s,
                                         ArcIteratorData<A> *data) const {
  GetImpl()->SetInstFuture(s, &inst_);
  GetImpl()->SetInstNode(&inst_);
  data->base = std::make_unique<ArcIterator<NGramFst<A>>>(*this, s);
}

namespace internal {

template <typename A>
NGramFstImpl<A>::NGramFstImpl(const Fst<A> &fst,
                              std::vector<StateId> *order_out) {
  typedef A Arc;
  typedef typename Arc::Label Label;
  typedef typename Arc::Weight Weight;
  typedef typename Arc::StateId StateId;
  SetType("ngram");
  SetInputSymbols(fst.InputSymbols());
  SetOutputSymbols(fst.OutputSymbols());
  SetProperties(kStaticProperties);

  // Check basic requirements for an OpenGrm language model Fst.
  if (!NGramFst<A>::HasRequiredProps(fst)) {
    FSTERROR() << "NGramFst only accepts OpenGrm language models as input";
    SetProperties(kError, kError);
    return;
  }

  int64_t num_states = CountStates(fst);
  std::vector<Label> context(num_states, 0);

  // Find the unigram state by starting from the start state, following
  // epsilons.
  StateId unigram = fst.Start();
  while (true) {
    if (unigram == kNoStateId) {
      FSTERROR() << "Could not identify unigram state";
      SetProperties(kError, kError);
      return;
    }
    ArcIterator<Fst<A>> aiter(fst, unigram);
    if (aiter.Done()) {
      LOG(WARNING) << "Unigram state " << unigram << " has no arcs.";
      break;
    }
    if (aiter.Value().ilabel != 0) break;
    unigram = aiter.Value().nextstate;
  }

  // Each state's context is determined by the subtree it is under from the
  // unigram state.
  std::queue<std::pair<StateId, Label>> label_queue;
  std::vector<bool> visited(num_states);
  // Force an epsilon link to the start state.
  label_queue.push(std::make_pair(fst.Start(), 0));
  for (ArcIterator<Fst<A>> aiter(fst, unigram); !aiter.Done(); aiter.Next()) {
    label_queue.push(
        std::make_pair(aiter.Value().nextstate, aiter.Value().ilabel));
  }
  // investigate states in breadth first fashion to assign context words.
  while (!label_queue.empty()) {
    std::pair<StateId, Label> &now = label_queue.front();
    if (!visited[now.first]) {
      context[now.first] = now.second;
      visited[now.first] = true;
      for (ArcIterator<Fst<A>> aiter(fst, now.first); !aiter.Done();
           aiter.Next()) {
        const Arc &arc = aiter.Value();
        if (arc.ilabel != 0) {
          label_queue.push(std::make_pair(arc.nextstate, now.second));
        }
      }
    }
    label_queue.pop();
  }
  visited.clear();

  // The arc from the start state should be assigned an epsilon to put it
  // in front of the all other labels (which makes Start state 1 after
  // unigram which is state 0).
  context[fst.Start()] = 0;

  // Build the tree of contexts fst by reversing the epsilon arcs from fst.
  VectorFst<Arc> context_fst;
  uint64_t num_final = 0;
  for (int i = 0; i < num_states; ++i) {
    if (fst.Final(i) != Weight::Zero()) {
      ++num_final;
    }
    context_fst.SetFinal(context_fst.AddState(), fst.Final(i));
  }
  context_fst.SetStart(unigram);
  context_fst.SetInputSymbols(fst.InputSymbols());
  context_fst.SetOutputSymbols(fst.OutputSymbols());
  int64_t num_context_arcs = 0;
  int64_t num_futures = 0;
  for (StateIterator<Fst<A>> siter(fst); !siter.Done(); siter.Next()) {
    const StateId &state = siter.Value();
    num_futures += fst.NumArcs(state) - fst.NumInputEpsilons(state);
    ArcIterator<Fst<A>> aiter(fst, state);
    if (!aiter.Done()) {
      const Arc &arc = aiter.Value();
      // this arc goes from state to arc.nextstate, so create an arc from
      // arc.nextstate to state to reverse it.
      if (arc.ilabel == 0) {
        context_fst.AddArc(arc.nextstate, Arc(context[state], context[state],
                                              arc.weight, state));
        num_context_arcs++;
      }
    }
  }
  if (num_context_arcs != context_fst.NumStates() - 1) {
    FSTERROR() << "Number of contexts arcs != number of states - 1";
    SetProperties(kError, kError);
    return;
  }
  if (context_fst.NumStates() != num_states) {
    FSTERROR() << "Number of contexts != number of states";
    SetProperties(kError, kError);
    return;
  }
  int64_t context_props =
      context_fst.Properties(kIDeterministic | kILabelSorted, true);
  if (!(context_props & kIDeterministic)) {
    FSTERROR() << "Input Fst is not structured properly";
    SetProperties(kError, kError);
    return;
  }
  if (!(context_props & kILabelSorted)) {
    ArcSort(&context_fst, ILabelCompare<Arc>());
  }

  uint64_t b64;
  Weight weight;
  Label label = kNoLabel;
  const size_t storage = Storage(num_states, num_futures, num_final);
  std::unique_ptr<MappedFile> data_region(MappedFile::Allocate(storage));
  char *data = static_cast<char *>(data_region->mutable_data());
  memset(data, 0, storage);
  size_t offset = 0;
  memcpy(data + offset, reinterpret_cast<char *>(&num_states),
         sizeof(num_states));
  offset += sizeof(num_states);
  memcpy(data + offset, reinterpret_cast<char *>(&num_futures),
         sizeof(num_futures));
  offset += sizeof(num_futures);
  memcpy(data + offset, reinterpret_cast<char *>(&num_final),
         sizeof(num_final));
  offset += sizeof(num_final);
  uint64_t *context_bits = reinterpret_cast<uint64_t *>(data + offset);
  offset += BitmapIndex::StorageSize(num_states * 2 + 1) * sizeof(b64);
  uint64_t *future_bits = reinterpret_cast<uint64_t *>(data + offset);
  offset +=
      BitmapIndex::StorageSize(num_futures + num_states + 1) * sizeof(b64);
  uint64_t *final_bits = reinterpret_cast<uint64_t *>(data + offset);
  offset += BitmapIndex::StorageSize(num_states) * sizeof(b64);
  Label *context_words = reinterpret_cast<Label *>(data + offset);
  offset += (num_states + 1) * sizeof(label);
  Label *future_words = reinterpret_cast<Label *>(data + offset);
  offset += num_futures * sizeof(label);
  offset = (offset + sizeof(weight) - 1) & ~(sizeof(weight) - 1);
  Weight *backoff = reinterpret_cast<Weight *>(data + offset);
  offset += (num_states + 1) * sizeof(weight);
  Weight *final_probs = reinterpret_cast<Weight *>(data + offset);
  offset += num_final * sizeof(weight);
  Weight *future_probs = reinterpret_cast<Weight *>(data + offset);
  int64_t context_arc = 0, future_arc = 0, context_bit = 0, future_bit = 0,
          final_bit = 0;

  // pseudo-root bits
  BitmapIndex::Set(context_bits, context_bit++);
  ++context_bit;
  context_words[context_arc] = label;
  backoff[context_arc] = Weight::Zero();
  context_arc++;

  ++future_bit;
  if (order_out) {
    order_out->clear();
    order_out->resize(num_states);
  }

  std::queue<StateId> context_q;
  context_q.push(context_fst.Start());
  StateId state_number = 0;
  while (!context_q.empty()) {
    const StateId &state = context_q.front();
    if (order_out) {
      (*order_out)[state] = state_number;
    }

    const Weight final_weight = context_fst.Final(state);
    if (final_weight != Weight::Zero()) {
      BitmapIndex::Set(final_bits, state_number);
      final_probs[final_bit] = final_weight;
      ++final_bit;
    }

    for (ArcIterator<VectorFst<A>> aiter(context_fst, state); !aiter.Done();
         aiter.Next()) {
      const Arc &arc = aiter.Value();
      context_words[context_arc] = arc.ilabel;
      backoff[context_arc] = arc.weight;
      ++context_arc;
      BitmapIndex::Set(context_bits, context_bit++);
      context_q.push(arc.nextstate);
    }
    ++context_bit;

    for (ArcIterator<Fst<A>> aiter(fst, state); !aiter.Done(); aiter.Next()) {
      const Arc &arc = aiter.Value();
      if (arc.ilabel != 0) {
        future_words[future_arc] = arc.ilabel;
        future_probs[future_arc] = arc.weight;
        ++future_arc;
        BitmapIndex::Set(future_bits, future_bit++);
      }
    }
    ++future_bit;
    ++state_number;
    context_q.pop();
  }

  if ((state_number != num_states) || (context_bit != num_states * 2 + 1) ||
      (context_arc != num_states) || (future_arc != num_futures) ||
      (future_bit != num_futures + num_states + 1) ||
      (final_bit != num_final)) {
    FSTERROR() << "Structure problems detected during construction";
    SetProperties(kError, kError);
    return;
  }

  Init(data, std::move(data_region));
}

template <typename A>
inline void NGramFstImpl<A>::Init(const char *data,
                                  std::unique_ptr<MappedFile> data_region) {
  data_region_ = std::move(data_region);
  data_ = data;
  size_t offset = 0;
  num_states_ = *(reinterpret_cast<const uint64_t *>(data_ + offset));
  offset += sizeof(num_states_);
  num_futures_ = *(reinterpret_cast<const uint64_t *>(data_ + offset));
  offset += sizeof(num_futures_);
  num_final_ = *(reinterpret_cast<const uint64_t *>(data_ + offset));
  offset += sizeof(num_final_);
  uint64_t bits;
  size_t context_bits = num_states_ * 2 + 1;
  size_t future_bits = num_futures_ + num_states_ + 1;
  context_ = reinterpret_cast<const uint64_t *>(data_ + offset);
  offset += BitmapIndex::StorageSize(context_bits) * sizeof(bits);
  future_ = reinterpret_cast<const uint64_t *>(data_ + offset);
  offset += BitmapIndex::StorageSize(future_bits) * sizeof(bits);
  final_ = reinterpret_cast<const uint64_t *>(data_ + offset);
  offset += BitmapIndex::StorageSize(num_states_) * sizeof(bits);
  context_words_ = reinterpret_cast<const Label *>(data_ + offset);
  offset += (num_states_ + 1) * sizeof(*context_words_);
  future_words_ = reinterpret_cast<const Label *>(data_ + offset);
  offset += num_futures_ * sizeof(*future_words_);
  offset = (offset + sizeof(*backoff_) - 1) & ~(sizeof(*backoff_) - 1);
  backoff_ = reinterpret_cast<const Weight *>(data_ + offset);
  offset += (num_states_ + 1) * sizeof(*backoff_);
  final_probs_ = reinterpret_cast<const Weight *>(data_ + offset);
  offset += num_final_ * sizeof(*final_probs_);
  future_probs_ = reinterpret_cast<const Weight *>(data_ + offset);

  context_index_.BuildIndex(context_, context_bits,
                            /*enable_select_0_index=*/true,
                            /*enable_select_1_index=*/true);
  future_index_.BuildIndex(future_, future_bits,
                           /*enable_select_0_index=*/true,
                           /*enable_select_1_index=*/false);
  final_index_.BuildIndex(final_, num_states_);

  select_root_ = context_index_.Select0s(0);
  if (context_index_.Rank1(0) != 0 || select_root_.first != 1 ||
      context_index_.Get(2) == false) {
    FSTERROR() << "Malformed file";
    SetProperties(kError, kError);
    return;
  }
  root_children_ = context_words_ + context_index_.Rank1(2);
  start_ = 1;
}

template <typename A>
inline typename A::StateId NGramFstImpl<A>::Transition(
    const std::vector<Label> &context, Label future) const {
  const Label *children = root_children_;
  size_t num_children = select_root_.second - 2;
  const Label *loc =
      std::lower_bound(children, children + num_children, future);
  if (loc == children + num_children || *loc != future) {
    return context_index_.Rank1(0);
  }
  size_t node = 2 + loc - children;
  size_t node_rank = context_index_.Rank1(node);
  std::pair<size_t, size_t> zeros =
      (node_rank == 0) ? select_root_ : context_index_.Select0s(node_rank);
  size_t first_child = zeros.first + 1;
  if (context_index_.Get(first_child) == false) {
    return node_rank;
  }
  size_t last_child = zeros.second - 1;
  for (int word = context.size() - 1; word >= 0; --word) {
    children = context_words_ + context_index_.Rank1(first_child);
    loc = std::lower_bound(children, children + last_child - first_child + 1,
                           context[word]);
    if (loc == children + last_child - first_child + 1 ||
        *loc != context[word]) {
      break;
    }
    node = first_child + loc - children;
    node_rank = context_index_.Rank1(node);
    zeros =
        (node_rank == 0) ? select_root_ : context_index_.Select0s(node_rank);
    first_child = zeros.first + 1;
    if (context_index_.Get(first_child) == false) break;
    last_child = zeros.second - 1;
  }
  return node_rank;
}

}  // namespace internal

/*****************************************************************************/
template <class A>
class NGramFstMatcher : public MatcherBase<A> {
 public:
  typedef A Arc;
  typedef typename A::Label Label;
  typedef typename A::StateId StateId;
  typedef typename A::Weight Weight;

  // This makes a copy of the FST.
  NGramFstMatcher(const NGramFst<A> &fst, MatchType match_type)
      : owned_fst_(fst.Copy()),
        fst_(*owned_fst_),
        inst_(fst_.inst_),
        match_type_(match_type),
        current_loop_(false),
        loop_(kNoLabel, 0, A::Weight::One(), kNoStateId) {
    if (match_type_ == MATCH_OUTPUT) {
      std::swap(loop_.ilabel, loop_.olabel);
    }
  }

  // This doesn't copy the FST.
  NGramFstMatcher(const NGramFst<A> *fst, MatchType match_type)
      : fst_(*fst),
        inst_(fst_.inst_),
        match_type_(match_type),
        current_loop_(false),
        loop_(kNoLabel, 0, A::Weight::One(), kNoStateId) {
    if (match_type_ == MATCH_OUTPUT) {
      std::swap(loop_.ilabel, loop_.olabel);
    }
  }

  // This makes a copy of the FST.
  NGramFstMatcher(const NGramFstMatcher<A> &matcher, bool safe = false)
      : owned_fst_(matcher.fst_.Copy(safe)),
        fst_(*owned_fst_),
        inst_(matcher.inst_),
        match_type_(matcher.match_type_),
        current_loop_(false),
        loop_(kNoLabel, 0, A::Weight::One(), kNoStateId) {
    if (match_type_ == MATCH_OUTPUT) {
      std::swap(loop_.ilabel, loop_.olabel);
    }
  }

  NGramFstMatcher<A> *Copy(bool safe = false) const override {
    return new NGramFstMatcher<A>(*this, safe);
  }

  MatchType Type(bool test) const override { return match_type_; }

  const Fst<A> &GetFst() const override { return fst_; }

  uint64_t Properties(uint64_t props) const override { return props; }

  void SetState(StateId s) final {
    fst_.GetImpl()->SetInstFuture(s, &inst_);
    current_loop_ = false;
  }

  bool Find(Label label) final {
    const Label nolabel = kNoLabel;
    done_ = true;
    if (label == 0 || label == nolabel) {
      if (label == 0) {
        current_loop_ = true;
        loop_.nextstate = inst_.state_;
      }
      // The unigram state has no epsilon arc.
      if (inst_.state_ != 0) {
        arc_.ilabel = arc_.olabel = 0;
        fst_.GetImpl()->SetInstNode(&inst_);
        arc_.nextstate = fst_.GetImpl()->context_index_.Rank1(
            fst_.GetImpl()->context_index_.Select1(
                fst_.GetImpl()->context_index_.Rank0(inst_.node_) - 1));
        arc_.weight = fst_.GetImpl()->backoff_[inst_.state_];
        done_ = false;
      }
    } else {
      current_loop_ = false;
      const Label *start = fst_.GetImpl()->future_words_ + inst_.offset_;
      const Label *end = start + inst_.num_futures_;
      const Label *search = std::lower_bound(start, end, label);
      if (search != end && *search == label) {
        size_t state = search - start;
        arc_.ilabel = arc_.olabel = label;
        arc_.weight = fst_.GetImpl()->future_probs_[inst_.offset_ + state];
        fst_.GetImpl()->SetInstContext(&inst_);
        arc_.nextstate = fst_.GetImpl()->Transition(inst_.context_, label);
        done_ = false;
      }
    }
    return !Done();
  }

  bool Done() const final { return !current_loop_ && done_; }

  const Arc &Value() const final { return (current_loop_) ? loop_ : arc_; }

  void Next() final {
    if (current_loop_) {
      current_loop_ = false;
    } else {
      done_ = true;
    }
  }

  ssize_t Priority(StateId s) final { return fst_.NumArcs(s); }

 private:
  std::unique_ptr<NGramFst<A>> owned_fst_;
  const NGramFst<A> &fst_;
  NGramFstInst<A> inst_;
  MatchType match_type_;  // Supplied by caller
  bool done_;
  Arc arc_;
  bool current_loop_;  // Current arc is the implicit loop
  Arc loop_;
};

/*****************************************************************************/
// Specialization for NGramFst; see generic version in fst.h
// for sample usage (but use the ProdLmFst type!). This version
// should inline.
template <class A>
class StateIterator<NGramFst<A>> : public StateIteratorBase<A> {
 public:
  typedef typename A::StateId StateId;

  explicit StateIterator(const NGramFst<A> &fst)
      : s_(0), num_states_(fst.NumStates()) {}

  bool Done() const final { return s_ >= num_states_; }

  StateId Value() const final {
    DCHECK(!Done());
    return s_;
  }

  void Next() final { ++s_; }

  void Reset() final { s_ = 0; }

 private:
  StateId s_;
  StateId num_states_;
};

/*****************************************************************************/
template <class A>
class ArcIterator<NGramFst<A>> : public ArcIteratorBase<A> {
 public:
  typedef A Arc;
  typedef typename A::Label Label;
  typedef typename A::StateId StateId;
  typedef typename A::Weight Weight;

  ArcIterator(const NGramFst<A> &fst, StateId state)
      : lazy_(~0), impl_(fst.GetImpl()), i_(0), flags_(kArcValueFlags) {
    inst_ = fst.inst_;
    impl_->SetInstFuture(state, &inst_);
    impl_->SetInstNode(&inst_);
  }

  bool Done() const final {
    return i_ >=
           ((inst_.node_ == 0) ? inst_.num_futures_ : inst_.num_futures_ + 1);
  }

  const Arc &Value() const final {
    DCHECK(!Done());
    bool eps = (inst_.node_ != 0 && i_ == 0);
    StateId state = (inst_.node_ == 0) ? i_ : i_ - 1;
    if (flags_ & lazy_ & (kArcILabelValue | kArcOLabelValue)) {
      arc_.ilabel = arc_.olabel =
          eps ? 0 : impl_->future_words_[inst_.offset_ + state];
      lazy_ &= ~(kArcILabelValue | kArcOLabelValue);
    }
    if (flags_ & lazy_ & kArcNextStateValue) {
      if (eps) {
        arc_.nextstate =
            impl_->context_index_.Rank1(impl_->context_index_.Select1(
                impl_->context_index_.Rank0(inst_.node_) - 1));
      } else {
        if (lazy_ & kArcNextStateValue) {
          impl_->SetInstContext(&inst_);  // first time only.
        }
        arc_.nextstate = impl_->Transition(
            inst_.context_, impl_->future_words_[inst_.offset_ + state]);
      }
      lazy_ &= ~kArcNextStateValue;
    }
    if (flags_ & lazy_ & kArcWeightValue) {
      arc_.weight = eps ? impl_->backoff_[inst_.state_]
                        : impl_->future_probs_[inst_.offset_ + state];
      lazy_ &= ~kArcWeightValue;
    }
    return arc_;
  }

  void Next() final {
    ++i_;
    lazy_ = ~0;
  }

  size_t Position() const final { return i_; }

  void Reset() final {
    i_ = 0;
    lazy_ = ~0;
  }

  void Seek(size_t a) final {
    if (i_ != a) {
      i_ = a;
      lazy_ = ~0;
    }
  }

  uint8_t Flags() const final { return flags_; }

  void SetFlags(uint8_t flags, uint8_t mask) final {
    flags_ &= ~mask;
    flags_ |= (flags & kArcValueFlags);
  }

 private:
  mutable Arc arc_;
  mutable uint8_t lazy_;
  const internal::NGramFstImpl<A> *impl_;  // Borrowed reference.
  mutable NGramFstInst<A> inst_;

  size_t i_;
  uint8_t flags_;
};

}  // namespace fst
#endif  // FST_EXTENSIONS_NGRAM_NGRAM_FST_H_