#ifndef FST_ARC_ARENA_H_
|
#define FST_ARC_ARENA_H_
|
|
#include <deque>
|
#include <memory>
|
#include <utility>
|
#include <fst/fst.h>
|
#include <fst/memory.h>
|
#include <unordered_map>
|
|
namespace fst {
|
|
// ArcArena is used for fast allocation of contiguous arrays of arcs.
|
//
|
// To create an arc array:
|
// for each state:
|
// for each arc:
|
// arena.PushArc();
|
// // Commits these arcs and returns pointer to them.
|
// Arc *arcs = arena.GetArcs();
|
//
|
// OR
|
//
|
// arena.DropArcs(); // Throws away current arcs, reuse the space.
|
//
|
// The arcs returned are guaranteed to be contiguous and the pointer returned
|
// will never be invalidated until the arena is cleared for reuse.
|
//
|
// The contents of the arena can be released with a call to arena.Clear() after
|
// which the arena will restart with an initial allocation capable of holding at
|
// least all of the arcs requested in the last usage before Clear() making
|
// subsequent uses of the Arena more efficient.
|
//
|
// The max_retained_size option can limit the amount of arc space requested on
|
// Clear() to avoid excess growth from intermittent high usage.
|
template <typename Arc>
|
class ArcArena {
|
public:
|
explicit ArcArena(size_t block_size = 256,
|
size_t max_retained_size = 1e6)
|
: block_size_(block_size),
|
max_retained_size_(max_retained_size) {
|
blocks_.emplace_back(MakeSharedBlock(block_size_));
|
first_block_size_ = block_size_;
|
total_size_ = block_size_;
|
arcs_ = blocks_.back().get();
|
end_ = arcs_ + block_size_;
|
next_ = arcs_;
|
}
|
|
ArcArena(const ArcArena& copy)
|
: arcs_(copy.arcs_), next_(copy.next_), end_(copy.end_),
|
block_size_(copy.block_size_),
|
first_block_size_(copy.first_block_size_),
|
total_size_(copy.total_size_),
|
max_retained_size_(copy.max_retained_size_),
|
blocks_(copy.blocks_) {
|
NewBlock(block_size_);
|
}
|
|
void ReserveArcs(size_t n) {
|
if (next_ + n < end_) return;
|
NewBlock(n);
|
}
|
|
void PushArc(const Arc& arc) {
|
if (next_ == end_) {
|
size_t length = next_ - arcs_;
|
NewBlock(length * 2);
|
}
|
*next_ = arc;
|
++next_;
|
}
|
|
const Arc* GetArcs() {
|
const auto *arcs = arcs_;
|
arcs_ = next_;
|
return arcs;
|
}
|
|
void DropArcs() { next_ = arcs_; }
|
|
size_t Size() { return total_size_; }
|
|
void Clear() {
|
blocks_.resize(1);
|
if (total_size_ > first_block_size_) {
|
first_block_size_ = std::min(max_retained_size_, total_size_);
|
blocks_.back() = MakeSharedBlock(first_block_size_);
|
}
|
total_size_ = first_block_size_;
|
arcs_ = blocks_.back().get();
|
end_ = arcs_ + first_block_size_;
|
next_ = arcs_;
|
}
|
|
private:
|
// Allocates a new block with capacity of at least n or block_size,
|
// copying incomplete arc sequence from old block to new block.
|
void NewBlock(size_t n) {
|
const auto length = next_ - arcs_;
|
const auto new_block_size = std::max(n, block_size_);
|
total_size_ += new_block_size;
|
blocks_.emplace_back(MakeSharedBlock(new_block_size));
|
std::copy(arcs_, next_, blocks_.back().get());
|
arcs_ = blocks_.back().get();
|
next_ = arcs_ + length;
|
end_ = arcs_ + new_block_size;
|
}
|
|
std::shared_ptr<Arc> MakeSharedBlock(size_t size) {
|
return std::shared_ptr<Arc>(new Arc[size], std::default_delete<Arc[]>());
|
}
|
|
Arc *arcs_;
|
Arc *next_;
|
const Arc *end_;
|
size_t block_size_;
|
size_t first_block_size_;
|
size_t total_size_;
|
size_t max_retained_size_;
|
std::list<std::shared_ptr<Arc>> blocks_;
|
};
|
|
// ArcArenaStateStore uses a resusable ArcArena to store arc arrays and does not
|
// require that the Expander call ReserveArcs first.
|
//
|
// TODO(tombagby): Make cache type configurable.
|
// TODO(tombagby): Provide ThreadLocal/Concurrent configuration.
|
template <class A>
|
class ArcArenaStateStore {
|
public:
|
using Arc = A;
|
using Weight = typename Arc::Weight;
|
using StateId = typename Arc::StateId;
|
|
ArcArenaStateStore() : arena_(64 * 1024) {
|
}
|
|
class State {
|
public:
|
Weight Final() const { return final_; }
|
|
size_t NumInputEpsilons() const { return niepsilons_; }
|
|
size_t NumOutputEpsilons() const { return noepsilons_; }
|
|
size_t NumArcs() const { return narcs_; }
|
|
const Arc &GetArc(size_t n) const { return arcs_[n]; }
|
|
const Arc *Arcs() const { return arcs_; }
|
|
int* MutableRefCount() const { return nullptr; }
|
|
private:
|
State(Weight weight, int32 niepsilons, int32 noepsilons, int32 narcs,
|
const Arc *arcs)
|
: final_(std::move(weight)),
|
niepsilons_(niepsilons),
|
noepsilons_(noepsilons),
|
narcs_(narcs),
|
arcs_(arcs) {}
|
|
Weight final_;
|
size_t niepsilons_;
|
size_t noepsilons_;
|
size_t narcs_;
|
const Arc *arcs_;
|
|
friend class ArcArenaStateStore<Arc>;
|
};
|
|
template <class Expander>
|
State *FindOrExpand(Expander &expander, StateId state_id) { // NOLINT
|
auto it = cache_.insert(std::pair<StateId, State*>(state_id, nullptr));
|
if (!it.second) return it.first->second;
|
// Needs a new state.
|
StateBuilder builder(&arena_);
|
expander.Expand(state_id, &builder);
|
const auto arcs = arena_.GetArcs();
|
size_t narcs = builder.narcs_;
|
size_t niepsilons = 0;
|
size_t noepsilons = 0;
|
for (size_t i = 0; i < narcs; ++i) {
|
if (arcs[i].ilabel == 0) ++niepsilons;
|
if (arcs[i].olabel == 0) ++noepsilons;
|
}
|
states_.emplace_back(
|
State(builder.final_, niepsilons, noepsilons, narcs, arcs));
|
// Places it in the cache.
|
auto state = &states_.back();
|
it.first->second = state;
|
return state;
|
}
|
|
State *Find(StateId state_id) const {
|
auto it = cache_.find(state_id);
|
return (it == cache_.end()) ? nullptr : it->second;
|
}
|
|
private:
|
class StateBuilder {
|
public:
|
explicit StateBuilder(ArcArena<Arc>* arena)
|
: arena_(arena), final_(Weight::Zero()), narcs_(0) {}
|
|
void SetFinal(Weight weight) { final_ = std::move(weight); }
|
|
void ReserveArcs(size_t n) { arena_->ReserveArcs(n); }
|
|
void AddArc(const Arc &arc) {
|
++narcs_;
|
arena_->PushArc(arc);
|
}
|
|
private:
|
friend class ArcArenaStateStore<Arc>;
|
|
ArcArena<Arc> *arena_;
|
Weight final_;
|
size_t narcs_;
|
};
|
|
std::unordered_map<StateId, State *> cache_;
|
std::deque<State> states_;
|
ArcArena<Arc> arena_;
|
};
|
|
} // namespace fst
|
|
#endif // FST_ARC_ARENA_H_
|
