compose-lattice-pruned.cc

Go to the documentation of this file.

// lat/compose-lattice-pruned.cc

// Copyright 2009-2012  Microsoft Corporation
//           2012-2013  Johns Hopkins University (Author: Daniel Povey)
//                2014  Guoguo Chen

// See ../../COPYING for clarification regarding multiple authors
//
// 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
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include "lat/compose-lattice-pruned.h"
#include "lat/lattice-functions.h"

namespace kaldi {

class PrunedCompactLatticeComposer {
 public:
  PrunedCompactLatticeComposer(
      const ComposeLatticePrunedOptions &opts,
      const CompactLattice &clat,
      fst::DeterministicOnDemandFst<fst::StdArc> *det_fst,
      CompactLattice* composed_clat);

  // Does the composition.  You must call this just once per object.
  void Compose();

 private:

  // Gets the num-arcs limit for this iteration of the algorithm, which will be
  // opts_.initial_num_arcs if there are currently no arcs; or otherwise
  // opts_.growth_ration * the current number of arcs (subject to the
  // opts_.max_arcs limit if we have already reached a final-state).  This helps
  // ensure that we call RecomputePruningInfo() on an appropriate schedule.
  int32 GetCurrentArcLimit() const;

  // This function, called just once at the start, computes all the static
  // information about the input lattice 'clat', in lat_state_info_.  (however,
  // the 'composed_states' members are just set to the empty vector for now.
  void ComputeLatticeStateInfo();

  // Called just once at the start, this sets up the first state in the
  // composed output.
  void AddFirstState();

  // This function processes the next un-expanded transition (or final-state)
  // out of the composed state numbered 'composed_state_to_expand'.
  void ProcessQueueElement(int32 composed_state_to_expand);

  // This is a part of ProcessQueueElements() that has been broken out
  // for clarity. it process the arc_index'th arc out of this source state.
  void ProcessTransition(int32 composed_src_state,
                         int32 arc_index);

  // This function recomputes certain members of the ComposedStateInfo relating
  // to the output states: namely, 'forward_cost', 'backward_cost' and
  // 'delta_backward_cost'.  In between calls to this function, we try to
  // keep those quantities as accurate as possible, but they aren't
  // completely accurate (see comments by their declarations for more info).
  void RecomputePruningInfo();

  // Sets '*composed_states' to a list of the states that currently
  // exist in the composed output, in topologically sorted order.
  // At exit, *composed_states will be a permutation of numbers
  // [0, 1, ...  clat_out_->NumStates() - 1], beginning with the
  // start-state 0.
  void GetTopsortedStateList(std::vector<int32> *composed_states) const;

  // Called from RecomputePruningInfo(), this computes all the 'forward_cost'
  // and 'prev_composed_state' members of the ComposedStateInfo.
  //   @param [in] composed_states  This is expected to be a list,
  //         in topological order, of all currently existing composed states,
  //         as produced by GetTopsortedStateList().
  void ComputeForwardCosts(const std::vector<int32> &composed_states);

  // Called from RecomputePruningInfo(), this computes all the 'backward_cost'
  // members of the ComposedStateInfo.  It also sets 'output_best_cost_'.
  // 'composed_states' is expected to be a list, in topological order, of all
  // currently existing composed states, as produced by GetTopsortedStateList().
  void ComputeBackwardCosts(const std::vector<int32> &composed_states);

  // Called from RecomputePruningInfo(), this computes all the
  // 'delta_backward_cost' members of the ComposedStateInfo.  'composed_states'
  // is expected to be a list, in topological order, of all currently existing
  // composed states, as produced by GetTopsortedStateList().  It also computes
  // the 'expected_cost_offset' values for all states, and uses them recreate
  // 'composed_state_queue_'.
  void ComputeDeltaBackwardCosts(const std::vector<int32> &composed_states);


  // This struct contains information about a state of the input lattice.
  struct LatticeStateInfo {
    // 'backward_cost' is the total cost of the best path from this state to
    //  the final state in the source lattice, including the final-prob.
    double backward_cost;

    // 'arc_delta_costs' is an array, one for each arc (and the final-prob, if
    // present), showing how much the cost to the final-state for the best path
    // starting in this state and exiting through each arc (or final-prob),
    // differs from 'backward_cost'.  Specifically, it contains pairs
    // (delta_cost, arc_index), where delta_cost >= 0 and arc_index is
    // either the index into this state's array of arcs (for arcs), or -1
    // if this represents the final-prob.
    //
    // 'arc_delta_costs' will be sorted, so that the first element has
    // .first=0.0 and the delta-costs will be increasing order.  This means that
    // we expand them from the start of the array, in order to process the best
    // arcs first.
    // lat_state_info_[i].arc_delta_costs.size() will equal will equal
    // clat_.NumStates(i), plus one if clat_.Final(i) != Zero().
    std::vector<std::pair<BaseFloat, int32> > arc_delta_costs;


    // 'composed_states' is a list of the state-ids in the composed output
    // that correspond to this state in the lattice, so we expect
    // that composed_state_info_[composed_states[i]].lat_state
    // equals the index of this lattice state.  This is helpful in
    // accessing the states in the output lattice in topological
    // order.
    std::vector<int32> composed_states;
  };

  // This struct contains information about a state of the composed
  // output.
  struct ComposedStateInfo {
    // 'lat_state' and 'lm_state' form the pair of states in the two FSTs
    // that this state corresponds to.  The unordered map 'pair_to_state_' maps these
    // state-pairs to the index of the composed state (the state-index in clat_out_).
    int32 lat_state;
    int32 lm_state;

    // the number of arcs on the path from the start state to this state, in the
    // composed lattice, by which this state was first reached.
    int32 depth;

    // If you have just called RecomputePruningInfo(), then
    // 'forward_cost' will equal the cost of the best path from the start-state
    // to this state, in the composed output.
    //
    // In between calls to RecomputePruningInfo() it may not always be fully up
    // to date; instead it will be an upper bound on what it would be if you had
    // just called RecomputePruningInfo(); it will be the cost of some path but
    // not necessarily the best path.
    double forward_cost;

    // 'backward_cost' relates to the cost from this state to the final-state in
    // the composed FST.  (By this we mean, more precisely, the cost of the best
    // path from this state to any final state, including the final-prob in that
    // final state).
    //
    // If we have just called RecomputePruningInfo(), then the following rules
    // specify what the value of 'backward_cost' will be:
    //   - If a final state is reachable from this state, backward_cost
    //     will contain the cost of the best path from this state to the
    //     final state (including the corresponding final-prob).
    //   - Otherwise, it will contain +infinity.
    //
    // If RecomputePruningInfo() has not just been called), it may contain any
    // value that is >= the value the the rules above specify (since, for
    // existing states, we don't modify it between calls to
    // RecomputePruningInfo()).  For states that have been added since
    // RecomputePruningInfo() was last called, it will be infinity.
    double backward_cost;

    // 'delta_backward_cost' is a quantity that is used in our heuristic of the
    // cost to an end-state from expanding a previously un-expanded arc.  It is
    // an estimate of the difference between the backward cost in this struct
    // (this->backward_cost) and the backward cost in the input lattice
    // (LatticeStateInfo::backward_cost).  This term reflects the anticipated
    // extra costs from 'det_fst_', which, while fairly close to zero, may be
    // substantial enough to want to correct for.
    //
    // The following is the value that 'delta_backward_cost' will have if
    // RecomputePruningInfo() has just been called:
    //   - If backward_cost is finite (this state in the composed result can
    //    reach the final state via currently expanded states), then
    //    delta_backward_cost is this->backward_cost minus
    //    lat_state_info_[lat_state].backward_cost.  (It will mostly, but
    //    not always, be <= 0, reflecting that the new LM is better than
    //    the old LM).
    //  - On the other hand, if backward_cost is infinite: delta_backward_cost
    //     is set to the delta_backward_cost of the previous state on the best
    //     path from the start state of the composed result to this state (or
    //     zero if this is the start state).
    //
    // If RecomputePruningInfo() has not just been called, then:
    //  - For states created since RecomputePruningInfo() was last called,
    //    delta_backward_cost will be inherited from the source state from
    //    which the new state was expanded.
    //  - For other states, delta_backward_cost will be unchanged since
    //    RecomputePruningInfo() was last called.
    // The above rules may make the delta_backward_cost a less accurate, but
    // still probably reasonable, heuristic.  What it is a heuristic for,
    // is: if we were to successfully reach an end-state of the composed output
    // from this state, what would be the resulting backward_cost
    // minus lat_state_info_[lat_state].backward_cost.
    BaseFloat delta_backward_cost;

    // 'prev_composed_state' is the previous state on the best path from
    // the start-state to the current state (or -1 if this is the start state).
    // It is computed in RecomputePruningInfo() when setting up 'forward_cost',
    // and then used to compute delta_backward_cost.  It is not otherwise
    // used.
    int32 prev_composed_state;

    // 'sorted_arc_index' is an index into the 'arc_delta_costs' array which is
    // a member of the LatticeStateInfo object corresponding to the lattice
    // state 'lat_state'.  It corresponds to the next arc (or final-prob) out of
    // the input lattice that we have yet to expand in the composition; or -1 if
    // we have expanded all of them.  When we first reach a composed state,
    // 'sorted_arc_index' will be zero; then it will increase one at a time as
    // we expand arcs until either the composition terminates or we have
    // expanded all the arcs and it becomes -1.
    int32 sorted_arc_index;

    // 'arc_delta_cost' is a derived quantity that we store here for easier
    // access.  Suppose this_lat_info is lat_state_info_[lat_state]; then
    // if sorted_arc_index >= 0, then:
    //    arc_delta_cost == this_lat_info.arc_delta_costs[sorted_arc_index].first
    // else: arc_delta_cost == +infinity.
    //
    // what 'arc_delta_cost' represents (or is a heuristic for), is the expected
    // cost of a path to the final-state leaving through the arc we're about to
    // expand, minus the expected cost of any path to the final-state starting
    // from this state.
    BaseFloat arc_delta_cost;

    // view 'expected_cost_offset' a phantom field of this struct, that has
    // been optimized out.  It's clearer if we act like it's a field, but
    // actually it's not stored.
    //
    // 'expected_cost_offset' is a derived quantity that reflects the expected
    // cost (according to our heuristic) of the best path we might encounter
    // when expanding the next previously unseen arc (or final-prob),
    // corresponding to 'sorted_arc_index'.  (This is the expected cost of a
    // successful path, from the beginning to the end of the lattice, but
    // constrained to be a path that contains the arc we're about to expand).
    //
    // The 'offset' part is about subtracting the best cost of the lattice, so we
    // can cast to float without too much loss of accuracy:
    //   expected_cost_offset = expected_cost - lat_best_cost_.
    //
    // We define expected_cost_offset by defining the 'expected_cost' part;
    // for clarity:
    //   First, let lat_backward_cost equal the backward_cost of the LatticeStateInfo
    //   corresponding to 'lat_state', i.e.
    //   lat_backward_cost = lat_state_info_[lat_state].backward_cost.  Then:
    //  expected_cost = forward_cost + lat_backward_cost +
    //                  delta_backward_cost + arc_delta_cost.
    // expected_cost_offset will always equal the above minus lat_best_cost_.
    //
    // The formula for expected_cost above is a pretty good heuristic for what
    // the cost to the end-state will be.  If the costs in det_fst_ were zero,
    // then the expression (forward_cost + lat_backward_cost + arc_delta_cost)
    // would be exact, and we would expand things in the ideal, best-first
    // order.  "delta_backward_cost" is a reasonable approximation for the extra
    // costs from 'det_fst_'.
    // BaseFloat expected_cost_offset;
  };

  // This bool variable is initialized to false, and will be updated to true
  // the first time a Final() function is called on the det_fst_. Then we will
  // immediately call RecomputeRruningInfo() so that the output_best_cost_ is
  // changed from +inf to a finite value, to be used in beam search. This is the
  // only time the RecomputeRruningInfo() function is called manually; otherwise
  // it always follows an automatic schedule based on the num-arcs of the output
  // lattice.
  bool output_reached_final_;

  // This variable, which we set initially to -1000, makes sure that in the
  // beginning of the algorithm, we always prioritize exploring the lattice
  // in a depth-first way. Once we find a path reaching a final state, this
  // variable will be reset to 0.
  // The reason we do this is because the beam-search depends on a good estimate
  // of the composed-best-cost, which before we reach a final state, we instead
  // borrow the value from best-cost from the input lattice, which is usually
  // systematically worse than the RNNLM scores, and makes the algorithm spend
  // a lot of time before reaching any final state, especially if the input
  // lattices are large.
  float depth_penalty_;
  const ComposeLatticePrunedOptions &opts_;
  const CompactLattice &clat_in_;
  fst::DeterministicOnDemandFst<fst::StdArc> *det_fst_;
  CompactLattice *clat_out_;

  // This counter keeps track of the number of arcs in the output lattice
  // clat_out_.  When it exceeds max_arcs,
  int32 num_arcs_out_;

  std::vector<LatticeStateInfo> lat_state_info_;

  // 'lat_best_cost' is the cost of the best path in the input lattice,
  // equal to lat_state_info_[0].backward_cost (we check that 0 is the
  // start state in the input lattice).
  double lat_best_cost_;

  // 'output_best_cost_' is the cost of the best successful path in the output
  // lattice 'clat_out_'; or +infinity if 'clat_out_' does not yet have any
  // successful paths.  It is updated only when RecomputePruningInfo() is
  // called.
  double output_best_cost_;


  // current_cutoff_ is a value used in deciding which composed states
  // need to be included in the queue.  Each time RecomputePruningInfo()
  // called, current_cutoff_ is set to
  //    (output_best_cost_ - lat_best_cost_ + opts_.lattice_compose_beam).
  // It will be +infinity if the output lattice doesn't yet have any
  // successful paths.  It decreases with time.  You can compare the
  // phantom 'expected_cost_offset' members of ComposedStateInfo with this
  // value; if they are more than this value, then there is no need
  // to enter the corresponding state into the queue.
  BaseFloat current_cutoff_;

  typedef std::priority_queue<std::pair<BaseFloat, int32>,
                      std::vector<std::pair<BaseFloat, int32> >,
                      std::greater<std::pair<BaseFloat, int32> > > QueueType;

  // composed_state_queue_ is a priority queue of the composed states
  // that we are intending to expand.  It contains pairs
  //   (expected_cost_offset, composed_state_index),
  // where expected_cost_offset == the phantom variable
  //       composed_state_info_[composed_state_index].expected_cost_offset.
  // We process the states from lowest cost first.
  // Every time RecomputePruningInfo() is called, this is cleared and repopulated
  // (since the states' expected_cost_offset values may change), and in between
  // calls to RecomputePruningInfo(), we do insert elements for newly created
  // states.
  QueueType composed_state_queue_;


  std::vector<ComposedStateInfo> composed_state_info_;

  // This maps a pair (lat_state, lm_state) to the index of the
  // state in the composed FST.  That would correspond to a state-id in
  // clat_out_, and also to an index into 'composed_state_info_'.
  unordered_map<std::pair<int32,int32>,
                int32, PairHasher<int32> > pair_to_state_;

  // This contains the set of state-indexes of the input lattice that already
  // have states in the composed output (i.e. is in accessed_lat_states_ if and
  // only if !lat_state_info_[i].composed_states.empty().  The point is to be
  // able to enumerate, in order or in reverse order, just those states in the
  // lattice that appear in the composed output (it's an efficiency thing that
  // will matter more for early iterations of the composition, when we need
  // to access the output lattice in topological order).
  std::set<int32> accessed_lat_states_;
};


void PrunedCompactLatticeComposer::GetTopsortedStateList(
    std::vector<int32> *composed_states) const {
  composed_states->clear();
  composed_states->reserve(clat_out_->NumStates());
  std::set<int32>::const_iterator iter = accessed_lat_states_.begin(),
      end = accessed_lat_states_.end();
  for (; iter != end; ++iter) {
    int32 lat_state = *iter;
    const LatticeStateInfo &input_lat_info = lat_state_info_[lat_state];
    composed_states->insert(composed_states->end(),
                            input_lat_info.composed_states.begin(),
                            input_lat_info.composed_states.end());
  }
  KALDI_ASSERT((*composed_states)[0] == 0 &&
               static_cast<int32>(composed_states->size()) ==
               clat_out_->NumStates());
}

int32 PrunedCompactLatticeComposer::GetCurrentArcLimit() const {
  int32 current_num_arcs = num_arcs_out_;
  if (current_num_arcs == 0) {
    return opts_.initial_num_arcs;
  } else {
    KALDI_ASSERT(opts_.growth_ratio > 1.0);
    int32 ans = static_cast<int32>(current_num_arcs *
                                   opts_.growth_ratio);
    if (ans == current_num_arcs)  // make sure the target increases.
      ans = current_num_arcs + 1;
    // if we have already reached a final state, then
    // apply the max_arcs limit.
    if (output_best_cost_ - output_best_cost_ == 0.0 &&
        ans > opts_.max_arcs)
      ans = opts_.max_arcs;
    return ans;
  }

}


void PrunedCompactLatticeComposer::RecomputePruningInfo() {
  std::vector<int32> all_composed_states;
  GetTopsortedStateList(&all_composed_states);
  ComputeForwardCosts(all_composed_states);
  ComputeBackwardCosts(all_composed_states);
  ComputeDeltaBackwardCosts(all_composed_states);
}

void PrunedCompactLatticeComposer::ComputeForwardCosts(
    const std::vector<int32> &composed_states) {
  KALDI_ASSERT(composed_states[0] == 0);

  // Note: when we initialized composed_state_info_[0]
  // we set forward_cost = 0.0, prev_composed_state = -1.

  std::vector<ComposedStateInfo>::iterator
      state_iter = composed_state_info_.begin(),
      state_end = composed_state_info_.end();

  state_iter->depth = 0;  // start state has depth 0
  ++state_iter;  // Skip over the start state.
  // Set all other forward_cost fields to infinity and prev_composed_state to
  // -1.
  for (; state_iter != state_end; ++state_iter) {
    state_iter->forward_cost = std::numeric_limits<double>::infinity();
    state_iter->prev_composed_state = -1;
  }

  std::vector<int32>::const_iterator state_index_iter = composed_states.begin(),
      state_index_end = composed_states.end();
  for (; state_index_iter != state_index_end; ++state_index_iter) {
    int32 composed_state_index = *state_index_iter;
    const ComposedStateInfo &info = composed_state_info_[
        composed_state_index];
    double forward_cost = info.forward_cost;
    // The next line is a check for infinity.  If infinities have appeared, it
    // either means there is a bug in the algorithm or there were infinities or
    // NaN's in the lattice.
    KALDI_ASSERT(forward_cost - forward_cost == 0.0);
    fst::ArcIterator<CompactLattice> aiter(*clat_out_,
                                           composed_state_index);
    for (; !aiter.Done(); aiter.Next()) {
      const CompactLatticeArc &arc = aiter.Value();
      double arc_cost = ConvertToCost(arc.weight),
          next_forward_cost = forward_cost + arc_cost;
      ComposedStateInfo &next_info = composed_state_info_[arc.nextstate];
      if (next_info.forward_cost > next_forward_cost) {
        next_info.forward_cost = next_forward_cost;
        next_info.prev_composed_state = composed_state_index;
        next_info.depth = composed_state_info_[composed_state_index].depth + 1;
      }
    }
  }
}

void PrunedCompactLatticeComposer::ComputeBackwardCosts(
    const std::vector<int32> &composed_states) {
  // Access the composed states in reverse topological order from latest to
  // earliest.
  std::vector<int32>::const_reverse_iterator iter = composed_states.rbegin(),
      end = composed_states.rend();
  for (; iter != end; ++iter) {
    int32 composed_state_index = *iter;
    ComposedStateInfo &info = composed_state_info_[composed_state_index];
    double backward_cost =
        ConvertToCost(clat_out_->Final(composed_state_index));
    fst::ArcIterator<CompactLattice> aiter(*clat_out_,
                                           composed_state_index);
    for (; !aiter.Done(); aiter.Next()) {
      const CompactLatticeArc &arc = aiter.Value();
      double arc_cost = ConvertToCost(arc.weight),
         next_backward_cost = composed_state_info_[arc.nextstate].backward_cost,
         this_backward_cost = arc_cost + next_backward_cost;
      if (this_backward_cost < backward_cost)
        backward_cost = this_backward_cost;
    }
    // It's OK if at this point, backward_cost is still +infinity.  This means
    // that this state cannot reach the end yet, which means we have not yet
    // expanded any path from this state all the way to a final-state of the
    // output.
    info.backward_cost = backward_cost;
  }
  output_best_cost_ = composed_state_info_[0].backward_cost;
  // See the declaration of current_cutoff_ for more information.  Note: on
  // early iterations, before any path reaches a final state of the composed
  // lattice, current_cutoff_ may be +infinity, and this is OK.
  current_cutoff_ =
      output_best_cost_ - lat_best_cost_ + opts_.lattice_compose_beam;
}

void PrunedCompactLatticeComposer::ComputeDeltaBackwardCosts(
    const std::vector<int32> &composed_states) {

  int32 num_states = clat_out_->NumStates();
  for (int32 composed_state_index = 0; composed_state_index < num_states;
       ++composed_state_index) {
    ComposedStateInfo &info = composed_state_info_[composed_state_index];
    int32 lat_state = info.lat_state;
    // Note: delta_backward_cost will be +infinity at this stage if the
    // backward_cost was +infinity.  This is OK; we'll set them all to
    // finite values later in this function.
    info.delta_backward_cost =
        info.backward_cost - lat_state_info_[lat_state].backward_cost + info.depth * depth_penalty_;
  }

  // 'queue_elements' is a list of items (expected_cost_offset,
  // composed_state_index) that we are going to add to composed_state_queue_,
  // after clearing it.  It's more efficient to accumulate them as a vector
  // and add them all at once, than adding them one by one (search online for
  // "heapify" if this seems confusing).
  std::vector<std::pair<BaseFloat, int32> > queue_elements;
  queue_elements.reserve(num_states);

  double lat_best_cost = lat_best_cost_;
  BaseFloat current_cutoff = current_cutoff_;
  std::vector<int32>::const_iterator iter = composed_states.begin(),
      end = composed_states.end();
  for (; iter != end; ++iter) {
    int32 composed_state_index = *iter;
    ComposedStateInfo &info = composed_state_info_[composed_state_index];
    if (info.delta_backward_cost - info.delta_backward_cost != 0) {
      // if info.delta_backward_cost is +infinity...
      int32 prev_composed_state = info.prev_composed_state;
      if (prev_composed_state < 0) {
        KALDI_ASSERT(composed_state_index == 0);
        info.delta_backward_cost = 0.0;
      } else {
        const ComposedStateInfo &prev_info =
            composed_state_info_[prev_composed_state];
        // Check that prev_info.delta_backward_cost is finite.
        KALDI_ASSERT(prev_info.delta_backward_cost -
                     prev_info.delta_backward_cost == 0.0);
        info.delta_backward_cost = prev_info.delta_backward_cost + depth_penalty_;
      }
    }
    double lat_backward_cost = lat_state_info_[info.lat_state].backward_cost;
    // See the formula by where expected_cost_offset is declared in the
    // struct for explanation.
    BaseFloat expected_cost_offset =
        info.forward_cost + lat_backward_cost + info.delta_backward_cost +
        info.arc_delta_cost - lat_best_cost;
    // If info.expected_cost_offset were real, we'd set it here:
    //info.expected_cost_offset = expected_cost_offset;

    // At this point expected_cost_offset may be infinite, if arc_delta_cost was
    // infinite (reflecting that we processed all the arcs, and the final-state
    // if applicable, of the lattice state corresponding to this composed state.
    if (expected_cost_offset < current_cutoff) {
      queue_elements.push_back(std::pair<BaseFloat, int32>(
          expected_cost_offset, composed_state_index));
    }
  }

  // Reinitialize composed_state_queue_ from 'queue_elements'.
  QueueType temp_queue(queue_elements.begin(), queue_elements.end());
  composed_state_queue_.swap(temp_queue);
}

void PrunedCompactLatticeComposer::ComputeLatticeStateInfo() {
  KALDI_ASSERT(clat_in_.Properties(fst::kTopSorted, true) ==
               fst::kTopSorted && clat_in_.NumStates() > 0 &&
               clat_in_.Start()  == 0);
  int32 num_lat_states = clat_in_.NumStates();
  lat_state_info_.resize(num_lat_states);

  for (int32 s = num_lat_states - 1; s >= 0; s--) {
    LatticeStateInfo &info = lat_state_info_[s];
    std::vector<std::pair<double, int32> > arc_costs;
    double backward_cost = ConvertToCost(clat_in_.Final(s));
    if (backward_cost != std::numeric_limits<double>::infinity())
      arc_costs.push_back(std::pair<BaseFloat,int32>(backward_cost, -1));
    fst::ArcIterator<CompactLattice> aiter(clat_in_, s);
    int32 arc_index = 0;
    for (; !aiter.Done(); aiter.Next(), ++arc_index)  {
      const CompactLatticeArc &arc = aiter.Value();
      KALDI_ASSERT(arc.nextstate > s);
      backward_cost = lat_state_info_[arc.nextstate].backward_cost +
          ConvertToCost(arc.weight);
      KALDI_ASSERT(backward_cost - backward_cost == 0.0 &&
                   "Possibly not all states of input lattice are co-accessible?");
      arc_costs.push_back(std::pair<BaseFloat,int32>(backward_cost, arc_index));
    }
    std::sort(arc_costs.begin(), arc_costs.end());
    KALDI_ASSERT(!arc_costs.empty() &&
                 "Possibly not all states of input lattice are co-accessible?");
    backward_cost = arc_costs[0].first;
    info.backward_cost = backward_cost;  // this is the state's backward_cost,
                                         // reflecting the best path to the end.
    info.arc_delta_costs.resize(arc_costs.size());
    std::vector<std::pair<double, int32> >::const_iterator
        src_iter = arc_costs.begin(), src_end = arc_costs.end();
    std::vector<std::pair<BaseFloat, int32> >::iterator
        dest_iter = info.arc_delta_costs.begin();
    for (; src_iter != src_end; ++src_iter, ++dest_iter) {
      dest_iter->first = BaseFloat(src_iter->first - backward_cost);
      dest_iter->second = src_iter->second;
    }
  }
  lat_best_cost_ = lat_state_info_[0].backward_cost;
}

PrunedCompactLatticeComposer::PrunedCompactLatticeComposer(
      const ComposeLatticePrunedOptions &opts,
      const CompactLattice &clat_in,
      fst::DeterministicOnDemandFst<fst::StdArc> *det_fst,
      CompactLattice* composed_clat): output_reached_final_(false),
    opts_(opts), clat_in_(clat_in), det_fst_(det_fst),
    clat_out_(composed_clat),
    num_arcs_out_(0),
    output_best_cost_(std::numeric_limits<double>::infinity()),
    current_cutoff_(std::numeric_limits<double>::infinity()) {
  clat_out_->DeleteStates();
  depth_penalty_ = -1000;
}


void PrunedCompactLatticeComposer::AddFirstState() {
  int32 state_id = clat_out_->AddState();
  clat_out_->SetStart(state_id);
  KALDI_ASSERT(state_id == 0);
  composed_state_info_.resize(1);
  ComposedStateInfo &composed_state = composed_state_info_[0];
  composed_state.lat_state = 0;
  composed_state.lm_state = det_fst_->Start();
  composed_state.depth = 0;
  composed_state.forward_cost = 0.0;
  composed_state.backward_cost = std::numeric_limits<double>::infinity();
  composed_state.delta_backward_cost = 0.0;
  composed_state.prev_composed_state = -1;
  composed_state.sorted_arc_index = 0;
  composed_state.arc_delta_cost = 0.0; // the first arc_delta_cost is always 0.0
                                       // due to sorting; no need to look it up.
  lat_state_info_[0].composed_states.push_back(state_id);
  accessed_lat_states_.insert(state_id);
  pair_to_state_[std::pair<int32, int32>(0, det_fst_->Start())] = state_id;

  BaseFloat expected_cost_offset = 0.0;  // the formula simplifies to zero
                                         // in this case.
  composed_state_queue_.push(
      std::pair<BaseFloat, int32>(expected_cost_offset,
                                  state_id));  // actually (0.0, 0).

}


void PrunedCompactLatticeComposer::ProcessQueueElement(
    int32 src_composed_state) {
  KALDI_ASSERT(static_cast<size_t>(src_composed_state) <
               composed_state_info_.size());

  ComposedStateInfo &src_composed_state_info = composed_state_info_[
      src_composed_state];
  int32 lat_state = src_composed_state_info.lat_state;
  const LatticeStateInfo &lat_state_info =
      lat_state_info_[lat_state];

  int32 sorted_arc_index = src_composed_state_info.sorted_arc_index,
      num_sorted_arcs = lat_state_info.arc_delta_costs.size();
  // note: num_sorted_arcs will be the number of arcs from this
  // lattice state; plus one if there is a final-prob.
  KALDI_ASSERT(sorted_arc_index >= 0);

  { // this block update the state's 'sorted_arc_index', 'arc_delta_cost' and
    // 'expected_cost_offset' to reflect the fact that (by the time we exit from
    // this function) we will have processed this arc (or the final-prob);
    // it also re-inserts this state into the queue, if appropriate.
    BaseFloat expected_cost_offset;
    if (sorted_arc_index + 1 == num_sorted_arcs) {
      src_composed_state_info.sorted_arc_index = -1;
      src_composed_state_info.arc_delta_cost =
          std::numeric_limits<BaseFloat>::infinity();
      expected_cost_offset =
          std::numeric_limits<BaseFloat>::infinity();
    } else {
      src_composed_state_info.sorted_arc_index = sorted_arc_index + 1;
      src_composed_state_info.arc_delta_cost =
          lat_state_info.arc_delta_costs[sorted_arc_index+1].first;
      expected_cost_offset =
          (src_composed_state_info.forward_cost +
           lat_state_info.backward_cost +
           src_composed_state_info.delta_backward_cost +
           src_composed_state_info.arc_delta_cost - lat_best_cost_);
    }
    // We do '<' here rather than '<=', so that if current_cutoff_ is infinity
    // and expected_cost_offset is infinity (because we've exhausted all the
    // transitions from this state, and sorted_arc_index is now -1), we don't
    // add this element to the queue.
    if (expected_cost_offset < current_cutoff_) {
      // this state has another exit arc (or final prob) that is good
      // enough to re-enter into the queue.  Note: if we are processing
      // an arc out of this state and the destination state is new,
      // we may also add something new to the queue at that time.

      // the following call should be equivalent to
      // composed_state_queue_.push(std::pair<BaseFloat,int32>(...)) with
      // the same pair of args.
      composed_state_queue_.emplace(
          expected_cost_offset, src_composed_state);
    }
  }

  int32 arc_index = lat_state_info.arc_delta_costs[sorted_arc_index].second;
  if (arc_index < 0) {  // This (arc_index == -1) means it is not really an arc
                        // index; it's a final-prob.
    int32 lm_state = src_composed_state_info.lm_state;
    BaseFloat lm_final_cost = det_fst_->Final(lm_state).Value();
    if (lm_final_cost != std::numeric_limits<BaseFloat>::infinity()) {
      // If there is a final-prob on this LM state (note: there always will be
      // for conventional language models), then add the final-prob of this
      // state...
      CompactLattice::Weight final_weight = clat_in_.Final(lat_state);
      // assume 'final_weight' is not Zero(); otherwise the final-prob should
      // not have been present in 'arc_delta_costs'.
      Lattice::Weight final_lat_weight = final_weight.Weight();
      final_lat_weight.SetValue1(final_lat_weight.Value1() +
                                 lm_final_cost);
      final_weight.SetWeight(final_lat_weight);
      clat_out_->SetFinal(src_composed_state, final_weight);
      double final_cost = ConvertToCost(final_lat_weight);
      if (final_cost < src_composed_state_info.backward_cost)
        src_composed_state_info.backward_cost = final_cost;
      if (!output_reached_final_) {
        output_reached_final_ = true;
        depth_penalty_ = 0.0;
        RecomputePruningInfo();
      }
    }
  } else {
    // It really was an arc.  This code is very complicated, so we make it its
    // own function.
    ProcessTransition(src_composed_state, arc_index);
  }
}

void PrunedCompactLatticeComposer::ProcessTransition(int32 src_composed_state,
                                                     int32 arc_index) {
  // Make src_composed_state a const pointer not a reference, as we may have to
  // modify the pointer if composed_state_info_ is resized.
  const ComposedStateInfo *src_info = &(composed_state_info_[
      src_composed_state]);
  int32 src_lat_state = src_info->lat_state;
  // Get the arc we are going to expand.
  fst::ArcIterator<CompactLattice> aiter(clat_in_, src_lat_state);
  aiter.Seek(arc_index);
  const CompactLatticeArc &lat_arc = aiter.Value();
  // Note: this code is for CompactLatticeArc, in which the ilabel and olabel
  // are the same, but we're writing it in such a way that it will naturally
  // generalize to LatticeArc, so there are separate variables for the ilabel
  // and the olabel.
  int32 dest_lat_state = lat_arc.nextstate,
      ilabel = lat_arc.ilabel,
      olabel = lat_arc.olabel;
  // Note: we expect that ilabel == olabel, since this is a CompactLattice, but this
  // may not be so if we extend this to work with Lattice.
  fst::StdArc lm_arc;

  // the input lattice might have epsilons
  if (olabel == 0) {
    lm_arc.ilabel = 0;
    lm_arc.olabel = 0;
    lm_arc.nextstate = src_info->lm_state;
    lm_arc.weight = fst::StdArc::Weight(0.0);
  } else if (!det_fst_->GetArc(src_info->lm_state, olabel, &lm_arc)) {
    // for normal language models we don't expect this to happen, but the
    // appropriate behavior is to do nothing; the composed arc does not exist,
    // so there is no arc to add and no new state to create.
    return;
  }
  int32 dest_lm_state = lm_arc.nextstate;
  // The following assertion is necessary because CompactLattice cannot support
  // different ilabel vs. olabel; and also it's an expectation about
  // language-models.
  KALDI_ASSERT(lm_arc.ilabel == lm_arc.olabel);

  LatticeStateInfo &dest_lat_state_info =
      lat_state_info_[dest_lat_state];

  int32 dest_composed_state;
  ComposedStateInfo *dest_info;

  { // The next block works out 'dest_composed_state' and
    // 'dest_info', and if the destination state did not already
    // exist, creates a new composed state.
    typedef std::unordered_map<std::pair<int32,int32>, int32,
        PairHasher<int32> > MapType;
    int32 new_composed_state = clat_out_->NumStates();
    std::pair<const std::pair<int32,int32>, int32> value(
        std::pair<int32,int32>(dest_lat_state, dest_lm_state), new_composed_state);
    std::pair<MapType::iterator, bool> ret =
        pair_to_state_.insert(value);
    if (ret.second) {
      // Successfully inserted: this dest-state did not already exist.  Most of
      // the rest of this block deals with the consequences of adding a new
      // state.
      int32 ans = clat_out_->AddState();
      KALDI_ASSERT(ans == new_composed_state);
      dest_composed_state = new_composed_state;
      composed_state_info_.resize(dest_composed_state + 1);
      dest_info = &(composed_state_info_[dest_composed_state]);
      // Re-assign src_composed_state as the vector might have been reallocated.
      src_info = &(composed_state_info_[src_composed_state]);
      if (dest_lat_state_info.composed_states.empty())
        accessed_lat_states_.insert(dest_lat_state);
      dest_lat_state_info.composed_states.push_back(new_composed_state);
      dest_info->lat_state = dest_lat_state;
      dest_info->lm_state = dest_lm_state;
      dest_info->depth = src_info->depth + 1;
      dest_info->forward_cost =
          src_info->forward_cost +
          ConvertToCost(lat_arc.weight) + lm_arc.weight.Value();
      dest_info->backward_cost =
          std::numeric_limits<double>::infinity();
      dest_info->delta_backward_cost =
          src_info->delta_backward_cost + dest_info->depth * depth_penalty_;
      // The 'prev_composed_state' field will not be read again until after it's
      // overwritten; we set it as below only for debugging purposes (the
      // negation is also for debugging purposes).
      dest_info->prev_composed_state = -src_composed_state;
      dest_info->sorted_arc_index = 0;
      dest_info->arc_delta_cost = 0.0;
      // Note: in the expression below, which can be understood with reference
      // to the comment by the declaration of the phantom variable
      // 'expected_cost_offset', 'arc_delta_cost' is known to equal 0.0 so it
      // has been removed.
      BaseFloat expected_cost_offset =
          (dest_info->forward_cost +
           dest_lat_state_info.backward_cost +
           dest_info->delta_backward_cost -
           lat_best_cost_);
      if (expected_cost_offset < current_cutoff_) {
        // the following call should be equivalent to
        // composed_state_queue_.push(std::pair<BaseFloat,int32>(...)) with
        // the same pair of args.
        composed_state_queue_.emplace(expected_cost_offset,
                                      dest_composed_state);
      }
    } else { // the destination composed state already existed.
      dest_composed_state = ret.first->second;
      dest_info = &(composed_state_info_[dest_composed_state]);
    }
  }
  // Add the arc from the src to dest state in the composed output.
  CompactLatticeArc new_arc;
  new_arc.nextstate = dest_composed_state;
  // Actually the ilabel and olabel are the same, but writing it this way will
  // generalize better to type Lattice if we need that later.
  new_arc.ilabel = ilabel;
  new_arc.olabel = olabel;
  new_arc.weight = lat_arc.weight;
  // 'weight' is the weight part, as opposed to the string part.
  LatticeArc::Weight weight = new_arc.weight.Weight();
  // include the LM-arc's weight in the weight of the new arc.
  weight.SetValue1(fst::Times(weight.Value1(), lm_arc.weight).Value());
  new_arc.weight.SetWeight(weight);
  clat_out_->AddArc(src_composed_state, new_arc);
  num_arcs_out_++;
}

static int32 TotalNumArcs(const CompactLattice &clat) {
  int32 num_states = clat.NumStates(),
      num_arcs = 0;
  for (int32 s = 0; s < num_states; s++)
    num_arcs += clat.NumArcs(s);
  return num_arcs;
}


void PrunedCompactLatticeComposer::Compose() {
  if (clat_in_.NumStates() == 0) {
    KALDI_WARN << "Input lattice to composition is empty.";
    return;
  }
  ComputeLatticeStateInfo();
  AddFirstState();
  // while (we have not reached final state  ||
  //        num-arcs produced < target num-arcs) { ...
  while (output_best_cost_ == std::numeric_limits<double>::infinity() ||
         num_arcs_out_ < opts_.max_arcs) {
    RecomputePruningInfo();
    int32 this_iter_arc_limit = GetCurrentArcLimit();
    while (num_arcs_out_ < this_iter_arc_limit &&
           !composed_state_queue_.empty()) {
      int32 src_composed_state = composed_state_queue_.top().second;
      composed_state_queue_.pop();
      ProcessQueueElement(src_composed_state);
    }
    if (composed_state_queue_.empty())
      break;
  }

  fst::Connect(clat_out_);
  TopSortCompactLatticeIfNeeded(clat_out_);

  if (GetVerboseLevel() >= 2) {
    int32 num_arcs_in = TotalNumArcs(clat_in_),
        orig_num_arcs_out = num_arcs_out_,
        num_arcs_out = TotalNumArcs(*clat_out_),
        num_states_in = clat_in_.NumStates(),
        orig_num_states_out = composed_state_info_.size(),
        num_states_out = clat_out_->NumStates();
    std::ostringstream os;
    os << "Input lattice had " << num_arcs_in << '/' << num_states_in
       << " arcs/states; output lattice has " << num_arcs_out << '/'
       << num_states_out;
    if (num_arcs_out != orig_num_arcs_out) {
      os << " (before pruning: " << orig_num_arcs_out << '/'
         << orig_num_states_out << ")";
    }
    if (!composed_state_queue_.empty()) {
      // Below, composed_state_queue_.top().first + lat_best_cost is an
      // expected-cost of the best path from the composed output that we *did
      // not* expand.  This, minus the best cost in the output compact lattice,
      // can be interpreted as the beam that we effecctively pruned the output
      // lattice to.
      BaseFloat effective_beam =
          composed_state_queue_.top().first + lat_best_cost_ - output_best_cost_;
      os << ". Effective beam was " << effective_beam;
    }
    KALDI_VLOG(2) << os.str();
  }

  if (clat_out_->NumStates() == 0) {
    KALDI_WARN << "Composed lattice has no states: something went wrong.";
  }
}

void ComposeCompactLatticePruned(
    const ComposeLatticePrunedOptions &opts,
    const CompactLattice &clat,
    fst::DeterministicOnDemandFst<fst::StdArc> *det_fst,
    CompactLattice* composed_clat) {
  PrunedCompactLatticeComposer composer(opts, clat, det_fst, composed_clat);
  composer.Compose();
}

} // namespace kaldi