use anyhow::Result;
use bitflags::bitflags;
use crate::algorithms::factor_weight::factor_iterators::{GallicFactorLeft, GallicFactorRight};
use crate::algorithms::factor_weight::{factor_weight, FactorWeightOptions, FactorWeightType};
use crate::algorithms::fst_convert::fst_convert_from_ref;
use crate::algorithms::tr_mappers::RmWeightMapper;
use crate::algorithms::weight_converters::{FromGallicConverter, ToGallicConverter};
use crate::algorithms::{reweight, shortest_distance, tr_map, weight_convert, ReweightType};
use crate::fst_impls::VectorFst;
use crate::fst_traits::{AllocableFst, ExpandedFst, MutableFst};
use crate::semirings::{DivideType, Semiring};
use crate::semirings::{
GallicWeightLeft, GallicWeightRight, StringWeightLeft, StringWeightRight,
WeaklyDivisibleSemiring, WeightQuantize,
};
bitflags! {
pub struct PushType: u32 {
const PUSH_WEIGHTS = 0b01;
const PUSH_LABELS = 0b10;
const REMOVE_TOTAL_WEIGHT = 0b100;
const REMOVE_COMMON_AFFIX = 0b1000;
}
}
pub fn push_weights<W, F>(
fst: &mut F,
reweight_type: ReweightType,
remove_total_weight: bool,
) -> Result<()>
where
F: MutableFst<W>,
W: WeaklyDivisibleSemiring,
{
let dist = shortest_distance(fst, reweight_type == ReweightType::ReweightToInitial)?;
if remove_total_weight {
let total_weight =
compute_total_weight(fst, &dist, reweight_type == ReweightType::ReweightToInitial)?;
reweight(fst, &dist, reweight_type)?;
remove_weight(
fst,
total_weight,
reweight_type == ReweightType::ReweightToFinal,
)?;
} else {
reweight(fst, &dist, reweight_type)?;
}
Ok(())
}
fn compute_total_weight<W, F>(fst: &F, dist: &[W], reverse: bool) -> Result<W>
where
W: Semiring,
F: ExpandedFst<W>,
{
if reverse {
if let Some(start) = fst.start() {
if start < dist.len() {
Ok(dist[start].clone())
} else {
Ok(W::zero())
}
} else {
Ok(W::zero())
}
} else {
let mut sum = W::zero();
for s in 0..dist.len() {
sum.plus_assign(
dist[s].times(unsafe { fst.final_weight_unchecked(s) }.unwrap_or_else(W::zero))?,
)?;
}
Ok(sum)
}
}
fn remove_weight<W, F>(fst: &mut F, weight: W, at_final: bool) -> Result<()>
where
F: MutableFst<W>,
W: WeaklyDivisibleSemiring,
{
if weight.is_one() || weight.is_zero() {
return Ok(());
}
if at_final {
unsafe {
for s in 0..fst.num_states() {
if let Some(mut final_weight) = fst.final_weight_unchecked(s) {
final_weight.divide_assign(&weight, DivideType::DivideRight)?;
fst.set_final_unchecked(s, final_weight);
}
}
}
} else if let Some(start) = fst.start() {
unsafe {
let mut it_tr = fst.tr_iter_unchecked_mut(start);
for idx_tr in 0..it_tr.len() {
let tr = it_tr.get_unchecked(idx_tr);
let weight = tr.weight.divide(&weight, DivideType::DivideLeft)?;
it_tr.set_weight_unchecked(idx_tr, weight);
}
if let Some(mut final_weight) = fst.final_weight_unchecked(start) {
final_weight.divide_assign(&weight, DivideType::DivideLeft)?;
fst.set_final_unchecked(start, final_weight);
}
}
}
Ok(())
}
macro_rules! m_labels_pushing {
($ifst: ident, $reweight_type: ident, $push_type: ident, $gallic_weight: ty, $string_weight: ident, $gallic_factor: ty) => {{
let mut mapper = ToGallicConverter {};
let mut gfst: VectorFst<$gallic_weight> = weight_convert($ifst, &mut mapper)?;
let gdistance = if $push_type.intersects(PushType::PUSH_WEIGHTS) {
shortest_distance(&gfst, $reweight_type == ReweightType::ReweightToInitial)?
} else {
let rm_weight_mapper = RmWeightMapper {};
let mut uwfst: VectorFst<_> = fst_convert_from_ref($ifst);
tr_map(&mut uwfst, &rm_weight_mapper)?;
let guwfst: VectorFst<$gallic_weight> = weight_convert(&uwfst, &mut mapper)?;
shortest_distance(&guwfst, $reweight_type == ReweightType::ReweightToInitial)?
};
if $push_type.intersects(PushType::REMOVE_COMMON_AFFIX | PushType::REMOVE_TOTAL_WEIGHT) {
let mut total_weight = compute_total_weight(
&gfst,
&gdistance,
$reweight_type == ReweightType::ReweightToInitial,
)?;
if !$push_type.intersects(PushType::REMOVE_COMMON_AFFIX) {
total_weight.set_value1($string_weight::one());
}
if !$push_type.intersects(PushType::REMOVE_TOTAL_WEIGHT) {
total_weight.set_value2(W::one());
}
reweight(&mut gfst, gdistance.as_slice(), $reweight_type)?;
remove_weight(
&mut gfst,
total_weight,
$reweight_type == ReweightType::ReweightToFinal,
)?;
} else {
reweight(&mut gfst, gdistance.as_slice(), $reweight_type)?;
}
let fwfst: VectorFst<$gallic_weight> =
factor_weight::<_, VectorFst<$gallic_weight>, _, _, $gallic_factor>(
&gfst,
FactorWeightOptions::new(
FactorWeightType::FACTOR_FINAL_WEIGHTS | FactorWeightType::FACTOR_ARC_WEIGHTS,
),
)?;
let mut mapper_from_gallic = FromGallicConverter {
superfinal_label: 0,
};
weight_convert(&fwfst, &mut mapper_from_gallic)
}};
}
pub fn push<W, F1, F2>(ifst: &F1, reweight_type: ReweightType, push_type: PushType) -> Result<F2>
where
F1: ExpandedFst<W>,
F2: ExpandedFst<W> + MutableFst<W> + AllocableFst<W>,
W: WeaklyDivisibleSemiring + WeightQuantize,
<W as Semiring>::ReverseWeight: 'static,
{
if push_type.intersects(PushType::PUSH_WEIGHTS) && !push_type.intersects(PushType::PUSH_LABELS)
{
let mut ofst = fst_convert_from_ref(ifst);
push_weights(
&mut ofst,
reweight_type,
push_type.intersects(PushType::REMOVE_TOTAL_WEIGHT),
)?;
Ok(ofst)
} else if push_type.intersects(PushType::PUSH_LABELS) {
match reweight_type {
ReweightType::ReweightToInitial => m_labels_pushing!(
ifst,
reweight_type,
push_type,
GallicWeightLeft<W>,
StringWeightLeft,
GallicFactorLeft<W>
),
ReweightType::ReweightToFinal => m_labels_pushing!(
ifst,
reweight_type,
push_type,
GallicWeightRight<W>,
StringWeightRight,
GallicFactorRight<W>
),
}
} else {
Ok(fst_convert_from_ref(ifst))
}
}