From ba7b512b9b0c56ea832f07b81cb977cb8db4e8e0 Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Fri, 16 Dec 2022 16:47:34 +0100
Subject: [PATCH 1/9] Lazy relabel fst

---
 Cargo.lock                                    |   6 +-
 rustfst/src/algorithms/compose/matcher_fst.rs |  25 +++
 rustfst/src/algorithms/mod.rs                 |   1 +
 rustfst/src/algorithms/relabel/mod.rs         |   2 +
 rustfst/src/algorithms/relabel/relabel_fst.rs | 172 ++++++++++++++++++
 .../src/algorithms/relabel/relabel_fst_op.rs  |  62 +++++++
 rustfst/src/algorithms/relabel_pairs.rs       |   2 +-
 .../src/tests_openfst/algorithms/compose.rs   |   1 +
 8 files changed, 267 insertions(+), 4 deletions(-)
 create mode 100644 rustfst/src/algorithms/relabel/mod.rs
 create mode 100644 rustfst/src/algorithms/relabel/relabel_fst.rs
 create mode 100644 rustfst/src/algorithms/relabel/relabel_fst_op.rs

diff --git a/Cargo.lock b/Cargo.lock
index 1427e2fdd..6d49a23e6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -696,7 +696,7 @@ dependencies = [
 
 [[package]]
 name = "rustfst"
-version = "0.13.0"
+version = "0.13.1"
 dependencies = [
  "anyhow",
  "bimap",
@@ -726,7 +726,7 @@ dependencies = [
 
 [[package]]
 name = "rustfst-cli"
-version = "0.13.0"
+version = "0.13.1"
 dependencies = [
  "anyhow",
  "clap",
@@ -740,7 +740,7 @@ dependencies = [
 
 [[package]]
 name = "rustfst-ffi"
-version = "0.13.0"
+version = "0.13.1"
 dependencies = [
  "anyhow",
  "downcast-rs",
diff --git a/rustfst/src/algorithms/compose/matcher_fst.rs b/rustfst/src/algorithms/compose/matcher_fst.rs
index abe2fad72..fcf9f7d5f 100644
--- a/rustfst/src/algorithms/compose/matcher_fst.rs
+++ b/rustfst/src/algorithms/compose/matcher_fst.rs
@@ -79,6 +79,7 @@ where
         let omatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchOutput)?;
 
         let mut add_on = (imatcher_data, omatcher_data);
+
         LabelLookAheadRelabeler::init(&mut fst, &mut add_on)?;
         LabelLookAheadRelabeler::relabel(fst2, &mut add_on, relabel_input)?;
 
@@ -91,6 +92,30 @@ where
             w: PhantomData,
         })
     }
+
+    // // Construct a new Matcher Fst intended for LookAhead composition and relabel fst2 wrt to the first fst.
+    // pub fn new_with_relabeling_2<F2: Fst<W>>(
+    //     mut fst: F,
+    //     fst2: &F2,
+    //     relabel_input: bool,
+    // ) -> Result<Self> {
+    //     let imatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchInput)?;
+    //     let omatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchOutput)?;
+    //
+    //     let mut add_on = (imatcher_data, omatcher_data);
+    //
+    //     LabelLookAheadRelabeler::init(&mut fst, &mut add_on)?;
+    //     LabelLookAheadRelabeler::relabel(fst2, &mut add_on, relabel_input)?;
+    //
+    //     let add_on = (add_on.0.map(Arc::new), add_on.1.map(Arc::new));
+    //
+    //     let fst_add_on = FstAddOn::new(fst, add_on);
+    //     Ok(Self {
+    //         fst_add_on,
+    //         matcher: PhantomData,
+    //         w: PhantomData,
+    //     })
+    // }
 }
 
 impl<W: Semiring, F: CoreFst<W>, B: Borrow<F>, M, T> CoreFst<W> for MatcherFst<W, F, B, M, T> {
diff --git a/rustfst/src/algorithms/mod.rs b/rustfst/src/algorithms/mod.rs
index 3d4fc45b1..d37ace4d6 100644
--- a/rustfst/src/algorithms/mod.rs
+++ b/rustfst/src/algorithms/mod.rs
@@ -50,6 +50,7 @@ mod partition;
 mod projection;
 mod push;
 mod queue;
+pub mod relabel;
 
 /// Module providing functions to randomly generate paths through an Fst. A static and a delayed version are available.
 pub mod randgen;
diff --git a/rustfst/src/algorithms/relabel/mod.rs b/rustfst/src/algorithms/relabel/mod.rs
new file mode 100644
index 000000000..142ad1260
--- /dev/null
+++ b/rustfst/src/algorithms/relabel/mod.rs
@@ -0,0 +1,2 @@
+mod relabel_fst;
+mod relabel_fst_op;
diff --git a/rustfst/src/algorithms/relabel/relabel_fst.rs b/rustfst/src/algorithms/relabel/relabel_fst.rs
new file mode 100644
index 000000000..964aa9d7d
--- /dev/null
+++ b/rustfst/src/algorithms/relabel/relabel_fst.rs
@@ -0,0 +1,172 @@
+use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache};
+use crate::fst_properties::FstProperties;
+use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
+use crate::{Semiring, StateId, SymbolTable, TrsVec};
+use anyhow::{Result, Context};
+use std::borrow::Borrow;
+use std::fmt::Debug;
+use std::sync::Arc;
+use crate::algorithms::relabel::relabel_fst_op::RelabelFstOp;
+use crate::algorithms::relabel_pairs::iterator_to_hashmap;
+
+type InnerLazyFst<W, F, B> = LazyFst<W, RelabelFstOp<W, F, B>, SimpleHashMapCache<W>>;
+
+struct RelabelFst<W: Semiring, F: Fst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RelabelFst<W, F, B> {
+    pub fn new<I, J>(fst: B, ipairs: I, opairs: J) -> Result<Self>
+    where
+        I: IntoIterator<Item = (StateId, StateId)>,
+        J: IntoIterator<Item = (StateId, StateId)>,
+    {
+        let map_ilabels = iterator_to_hashmap(ipairs)
+            .with_context(|| format_err!("Error while creating the HashMap for ipairs"))?;
+
+        let map_olabels = iterator_to_hashmap(opairs)
+            .with_context(|| format_err!("Error while creating the HashMap for opairs"))?;
+
+        let isymt = fst.borrow().input_symbols().cloned();
+        let osymt = fst.borrow().output_symbols().cloned();
+
+        let fst_op = RelabelFstOp::new(fst, map_ilabels, map_olabels);
+        let fst_cache = SimpleHashMapCache::default();
+        Ok(RelabelFst(LazyFst::from_op_and_cache(
+            fst_op, fst_cache, isymt, osymt,
+        )))
+    }
+
+    pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
+        self.0.compute()
+    }
+}
+
+impl<W, F, B> CoreFst<W> for RelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
+    type TRS = TrsVec<W>;
+
+    fn start(&self) -> Option<StateId> {
+        self.0.start()
+    }
+
+    fn final_weight(&self, state_id: StateId) -> Result<Option<W>> {
+        self.0.final_weight(state_id)
+    }
+
+    unsafe fn final_weight_unchecked(&self, state_id: StateId) -> Option<W> {
+        self.0.final_weight_unchecked(state_id)
+    }
+
+    fn num_trs(&self, s: StateId) -> Result<usize> {
+        self.0.num_trs(s)
+    }
+
+    unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
+        self.0.num_trs_unchecked(s)
+    }
+
+    fn get_trs(&self, state_id: StateId) -> Result<Self::TRS> {
+        self.0.get_trs(state_id)
+    }
+
+    unsafe fn get_trs_unchecked(&self, state_id: StateId) -> Self::TRS {
+        self.0.get_trs_unchecked(state_id)
+    }
+
+    fn properties(&self) -> FstProperties {
+        self.0.properties()
+    }
+
+    fn num_input_epsilons(&self, state: StateId) -> Result<usize> {
+        self.0.num_input_epsilons(state)
+    }
+
+    fn num_output_epsilons(&self, state: StateId) -> Result<usize> {
+        self.0.num_output_epsilons(state)
+    }
+}
+
+impl<'a, W, F, B> StateIterator<'a> for RelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'a,
+    B: Borrow<F> + 'a,
+{
+    type Iter = <InnerLazyFst<W, F, B> as StateIterator<'a>>::Iter;
+
+    fn states_iter(&'a self) -> Self::Iter {
+        self.0.states_iter()
+    }
+}
+
+impl<'a, W, F, B> FstIterator<'a, W> for RelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'a,
+    B: Borrow<F> + 'a,
+{
+    type FstIter = <InnerLazyFst<W, F, B> as FstIterator<'a, W>>::FstIter;
+
+    fn fst_iter(&'a self) -> Self::FstIter {
+        self.0.fst_iter()
+    }
+}
+
+impl<W, F, B> Fst<W> for RelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'static,
+    B: Borrow<F> + 'static,
+{
+    fn input_symbols(&self) -> Option<&Arc<SymbolTable>> {
+        self.0.input_symbols()
+    }
+
+    fn output_symbols(&self) -> Option<&Arc<SymbolTable>> {
+        self.0.output_symbols()
+    }
+
+    fn set_input_symbols(&mut self, symt: Arc<SymbolTable>) {
+        self.0.set_input_symbols(symt)
+    }
+
+    fn set_output_symbols(&mut self, symt: Arc<SymbolTable>) {
+        self.0.set_output_symbols(symt)
+    }
+
+    fn take_input_symbols(&mut self) -> Option<Arc<SymbolTable>> {
+        self.0.take_input_symbols()
+    }
+
+    fn take_output_symbols(&mut self) -> Option<Arc<SymbolTable>> {
+        self.0.take_output_symbols()
+    }
+}
+
+impl<W, F, B> Debug for RelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::fst_impls::VectorFst;
+    use crate::semirings::TropicalWeight;
+
+    use super::*;
+
+    #[test]
+    fn test_replace_fst_sync() {
+        fn is_sync<T: Sync>() {}
+        is_sync::<RelabelFst<TropicalWeight, VectorFst<_>, VectorFst<_>>>();
+    }
+}
diff --git a/rustfst/src/algorithms/relabel/relabel_fst_op.rs b/rustfst/src/algorithms/relabel/relabel_fst_op.rs
new file mode 100644
index 000000000..fc1931efb
--- /dev/null
+++ b/rustfst/src/algorithms/relabel/relabel_fst_op.rs
@@ -0,0 +1,62 @@
+use std::borrow::Borrow;
+use std::marker::PhantomData;
+
+use anyhow::Result;
+
+use crate::algorithms::lazy::FstOp;
+use crate::fst_properties::FstProperties;
+use crate::fst_traits::Fst;
+use crate::{Semiring, TrsVec, StateId, Trs};
+use std::fmt::{Debug, Formatter};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+pub struct RelabelFstOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
+    fst: B,
+    map_ilabels: HashMap<StateId, StateId>,
+    map_olabels: HashMap<StateId, StateId>,
+    ghost: PhantomData<(W, F)>,
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RelabelFstOp<W, F, B> {
+    pub fn new(fst: B, map_ilabels: HashMap<StateId, StateId>, map_olabels: HashMap<StateId, StateId>) -> Self {
+        Self {
+            fst, map_ilabels, map_olabels, ghost: PhantomData
+        }
+    }
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> Debug for RelabelFstOp<W, F, B> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        unimplemented!()
+    }
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> FstOp<W> for RelabelFstOp<W, F, B> {
+    fn compute_start(&self) -> Result<Option<u32>> {
+        Ok(self.fst.borrow().start())
+    }
+
+    fn compute_trs(&self, id: u32) -> Result<TrsVec<W>> {
+        let trs_original = self.fst.borrow().get_trs(id)?;
+        let mut trs = vec![];
+        for tr in trs_original.trs() {
+            let mut new_tr = tr.clone();
+            if let Some(new_ilabel) = self.map_ilabels.get(&tr.ilabel) {
+                new_tr.ilabel = *new_ilabel;
+            }
+            if let Some(new_olabel) = self.map_olabels.get(&tr.olabel) {
+                new_tr.olabel = *new_olabel;
+            }
+        }
+        Ok(TrsVec(Arc::new(trs)))
+    }
+
+    fn compute_final_weight(&self, id: u32) -> Result<Option<W>> {
+        self.fst.borrow().final_weight(id)
+    }
+
+    fn properties(&self) -> FstProperties {
+        unimplemented!()
+    }
+}
diff --git a/rustfst/src/algorithms/relabel_pairs.rs b/rustfst/src/algorithms/relabel_pairs.rs
index fd342ecef..722bf222d 100644
--- a/rustfst/src/algorithms/relabel_pairs.rs
+++ b/rustfst/src/algorithms/relabel_pairs.rs
@@ -7,7 +7,7 @@ use crate::fst_traits::MutableFst;
 use crate::semirings::Semiring;
 use crate::StateId;
 
-fn iterator_to_hashmap<I>(pairs: I) -> Result<HashMap<StateId, StateId>>
+pub fn iterator_to_hashmap<I>(pairs: I) -> Result<HashMap<StateId, StateId>>
 where
     I: IntoIterator<Item = (StateId, StateId)>,
 {
diff --git a/rustfst/src/tests_openfst/algorithms/compose.rs b/rustfst/src/tests_openfst/algorithms/compose.rs
index 6c2959e31..7615ab309 100644
--- a/rustfst/src/tests_openfst/algorithms/compose.rs
+++ b/rustfst/src/tests_openfst/algorithms/compose.rs
@@ -121,6 +121,7 @@ fn do_test_compose_lookahead<W>(
 where
     W: SerializableSemiring + WeightQuantize + WeaklyDivisibleSemiring,
 {
+    println!("Lookahead composition");
     type TLaFst<S, F> = MatcherFst<
         S,
         F,

From 1de3858853ffd8175c74b0f43ab1d2d763bfbd0f Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Mon, 19 Dec 2022 17:28:26 +0100
Subject: [PATCH 2/9] Implemement lazy lookahead relabeling

---
 .../src/algorithms/compose/label_reachable.rs |  42 +++++
 .../push_weights_compose_filter.rs            |   6 +-
 .../label_lookahead_relabeler.rs              |  21 ++-
 .../lookahead_relabel_fst.rs                  | 171 ++++++++++++++++++
 .../lookahead_relabel_fst_op.rs               |  77 ++++++++
 .../compose/lookahead_relabel/mod.rs          |   4 +
 rustfst/src/algorithms/compose/matcher_fst.rs |  49 ++---
 rustfst/src/algorithms/compose/mod.rs         |   3 +
 rustfst/src/algorithms/lazy/lazy_fst.rs       |  12 +-
 rustfst/src/algorithms/relabel/mod.rs         |   2 +
 rustfst/src/algorithms/relabel/relabel_fst.rs |   8 +-
 .../src/algorithms/relabel/relabel_fst_op.rs  |  24 ++-
 .../src/tests_openfst/algorithms/compose.rs   | 159 +++++++++++++++-
 13 files changed, 535 insertions(+), 43 deletions(-)
 create mode 100644 rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
 create mode 100644 rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
 create mode 100644 rustfst/src/algorithms/compose/lookahead_relabel/mod.rs

diff --git a/rustfst/src/algorithms/compose/label_reachable.rs b/rustfst/src/algorithms/compose/label_reachable.rs
index 12698419a..827b908b9 100644
--- a/rustfst/src/algorithms/compose/label_reachable.rs
+++ b/rustfst/src/algorithms/compose/label_reachable.rs
@@ -10,6 +10,7 @@ use crate::algorithms::{fst_convert_from_ref, tr_sort};
 use crate::fst_impls::VectorFst;
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::{CoreFst, ExpandedFst, Fst, MutableFst};
+use crate::prelude::compose::LookaheadRelabelFst;
 use crate::semirings::Semiring;
 use crate::{Label, StateId, Tr, Trs, EPS_LABEL, NO_LABEL, UNASSIGNED};
 
@@ -60,6 +61,19 @@ impl LabelReachableData {
             .or_insert_with(|| n as Label + 1)
     }
 
+    pub fn relabel_unmut(&self, label: Label) -> Result<Label> {
+        if label == EPS_LABEL {
+            return Ok(EPS_LABEL);
+        }
+        Ok(*self.label2index.get(&label).ok_or_else(|| {
+            format_err!(
+                "Missing label {:?} from relabelling table {:?}",
+                label,
+                self.label2index
+            )
+        })?)
+    }
+
     pub fn relabel_fst<W: Semiring, F: MutableFst<W>>(
         &mut self,
         fst: &mut F,
@@ -92,6 +106,34 @@ impl LabelReachableData {
         Ok(())
     }
 
+    pub fn relabel_fst_lazy<W: Semiring, F: Fst<W> + 'static>(
+        &mut self,
+        fst: Arc<F>,
+        relabel_input: bool,
+    ) -> Result<impl Fst<W>> {
+        // 1. Retrieve the relevant SymbolTable
+        let symt = if relabel_input {
+            fst.input_symbols()
+                .ok_or_else(|| format_err!("Input SymbolTable not attached to the fst which is mandatory for the relabel fst lazy to work"))?
+        } else {
+            fst.output_symbols()
+                .ok_or_else(|| format_err!("Output SymbolTable not attached to the fst which is mandatory for the relabel fst lazy to work"))?
+        };
+
+        // 2. Use the SymbolTable as a proxy to retrieve all the labels that are going to be ralabelled.
+        // By calling `relabel` on each of those we guarantee we won't see other labels at inference.
+        for (label, _) in symt.iter() {
+            self.relabel(label);
+        }
+
+        // 3. Create a custom Lazy Fst that will take as input the current LabelReachableData object
+        // with the whole relabeling mapping.
+        // This Lazy fst, will 1. use the mapping to relabel 2. sort the transitions
+        let la_relabel_fst = LookaheadRelabelFst::<W, F, _>::new(fst, self.clone(), relabel_input)?;
+
+        Ok(la_relabel_fst)
+    }
+
     // Returns relabeling pairs (cf. relabel.h::Relabel()). If avoid_collisions is
     // true, extra pairs are added to ensure no collisions when relabeling
     // automata that have labels unseen here.
diff --git a/rustfst/src/algorithms/compose/lookahead_filters/push_weights_compose_filter.rs b/rustfst/src/algorithms/compose/lookahead_filters/push_weights_compose_filter.rs
index 0b1b1d94d..f30b13824 100644
--- a/rustfst/src/algorithms/compose/lookahead_filters/push_weights_compose_filter.rs
+++ b/rustfst/src/algorithms/compose/lookahead_filters/push_weights_compose_filter.rs
@@ -12,7 +12,7 @@ use crate::algorithms::compose::lookahead_filters::LookAheadComposeFilterTrait;
 use crate::algorithms::compose::lookahead_matchers::{LookAheadMatcherData, LookaheadMatcher};
 use crate::algorithms::compose::matchers::{MatchType, MatcherFlags};
 use crate::fst_properties::FstProperties;
-use crate::fst_traits::{ExpandedFst, Fst};
+use crate::fst_traits::Fst;
 use crate::semirings::{
     DivideType, Semiring, SerializableSemiring, WeaklyDivisibleSemiring, WeightQuantize,
 };
@@ -80,8 +80,8 @@ impl<W, F1, F2, B1, B2, M1, M2, CFB, SMT> ComposeFilterBuilder<W, F1, F2, B1, B2
     for PushWeightsComposeFilterBuilder<W, F1, F2, B1, B2, M1, M2, CFB, SMT>
 where
     W: SerializableSemiring + WeaklyDivisibleSemiring + WeightQuantize,
-    F1: ExpandedFst<W>,
-    F2: ExpandedFst<W>,
+    F1: Fst<W>,
+    F2: Fst<W>,
     B1: Borrow<F1> + Debug,
     B2: Borrow<F2> + Debug,
     M1: LookaheadMatcher<W, F1, B1>,
diff --git a/rustfst/src/algorithms/compose/lookahead_matchers/label_lookahead_relabeler.rs b/rustfst/src/algorithms/compose/lookahead_matchers/label_lookahead_relabeler.rs
index 268f3efff..a682bef7b 100644
--- a/rustfst/src/algorithms/compose/lookahead_matchers/label_lookahead_relabeler.rs
+++ b/rustfst/src/algorithms/compose/lookahead_matchers/label_lookahead_relabeler.rs
@@ -1,8 +1,9 @@
 use anyhow::Result;
 
 use crate::algorithms::compose::LabelReachableData;
-use crate::fst_traits::MutableFst;
+use crate::fst_traits::{Fst, MutableFst};
 use crate::semirings::Semiring;
+use std::sync::Arc;
 
 pub struct LabelLookAheadRelabeler {}
 
@@ -39,4 +40,22 @@ impl LabelLookAheadRelabeler {
         }
         bail!("Addon contains only None elements")
     }
+
+    pub fn relabel_lazy<W: Semiring, F: Fst<W> + 'static>(
+        fst: Arc<F>,
+        addon: &mut (Option<LabelReachableData>, Option<LabelReachableData>),
+        relabel_input: bool,
+    ) -> Result<impl Fst<W>> {
+        if let Some(reachable_data) = &mut addon.0 {
+            let lazy_fst = reachable_data.relabel_fst_lazy(fst, relabel_input)?;
+            return Ok(lazy_fst);
+        }
+
+        if let Some(reachable_data) = &mut addon.1 {
+            let lazy_fst = reachable_data.relabel_fst_lazy(fst, relabel_input)?;
+            return Ok(lazy_fst);
+        }
+
+        bail!("Addon contains only None elements")
+    }
 }
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
new file mode 100644
index 000000000..a4006c7b6
--- /dev/null
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -0,0 +1,171 @@
+use crate::algorithms::compose::lookahead_relabel::lookahead_relabel_fst_op::LookaheadRelabelFstOp;
+use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache};
+use crate::fst_properties::FstProperties;
+use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
+use crate::prelude::compose::LabelReachableData;
+use crate::{Semiring, StateId, SymbolTable, TrsVec};
+use anyhow::Result;
+use std::borrow::Borrow;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+type InnerLazyFst<W, F, B> = LazyFst<W, LookaheadRelabelFstOp<W, F, B>, SimpleHashMapCache<W>>;
+
+pub struct LookaheadRelabelFst<W: Semiring, F: Fst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFst<W, F, B> {
+    pub fn new(
+        fst: B,
+        label_reachable_data: LabelReachableData,
+        relabel_input: bool,
+    ) -> Result<Self> {
+        let mut isymt = fst.borrow().input_symbols().cloned();
+        let mut osymt = fst.borrow().output_symbols().cloned();
+        if relabel_input {
+            isymt = None;
+        } else {
+            osymt = None;
+        }
+
+        let fst_op = LookaheadRelabelFstOp::new(fst, label_reachable_data, relabel_input);
+        let fst_cache = SimpleHashMapCache::default();
+        Ok(LookaheadRelabelFst(LazyFst::from_op_and_cache(
+            fst_op, fst_cache, isymt, osymt,
+        )))
+    }
+
+    pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
+        self.0.compute()
+    }
+}
+
+impl<W, F, B> CoreFst<W> for LookaheadRelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
+    type TRS = TrsVec<W>;
+
+    fn start(&self) -> Option<StateId> {
+        self.0.start()
+    }
+
+    fn final_weight(&self, state_id: StateId) -> Result<Option<W>> {
+        self.0.final_weight(state_id)
+    }
+
+    unsafe fn final_weight_unchecked(&self, state_id: StateId) -> Option<W> {
+        self.0.final_weight_unchecked(state_id)
+    }
+
+    fn num_trs(&self, s: StateId) -> Result<usize> {
+        self.0.num_trs(s)
+    }
+
+    unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
+        self.0.num_trs_unchecked(s)
+    }
+
+    fn get_trs(&self, state_id: StateId) -> Result<Self::TRS> {
+        self.0.get_trs(state_id)
+    }
+
+    unsafe fn get_trs_unchecked(&self, state_id: StateId) -> Self::TRS {
+        self.0.get_trs_unchecked(state_id)
+    }
+
+    fn properties(&self) -> FstProperties {
+        self.0.properties()
+    }
+
+    fn num_input_epsilons(&self, state: StateId) -> Result<usize> {
+        self.0.num_input_epsilons(state)
+    }
+
+    fn num_output_epsilons(&self, state: StateId) -> Result<usize> {
+        self.0.num_output_epsilons(state)
+    }
+}
+
+impl<'a, W, F, B> StateIterator<'a> for LookaheadRelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'a,
+    B: Borrow<F> + 'a,
+{
+    type Iter = <InnerLazyFst<W, F, B> as StateIterator<'a>>::Iter;
+
+    fn states_iter(&'a self) -> Self::Iter {
+        self.0.states_iter()
+    }
+}
+
+impl<'a, W, F, B> FstIterator<'a, W> for LookaheadRelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'a,
+    B: Borrow<F> + 'a,
+{
+    type FstIter = <InnerLazyFst<W, F, B> as FstIterator<'a, W>>::FstIter;
+
+    fn fst_iter(&'a self) -> Self::FstIter {
+        self.0.fst_iter()
+    }
+}
+
+impl<W, F, B> Fst<W> for LookaheadRelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W> + 'static,
+    B: Borrow<F> + 'static,
+{
+    fn input_symbols(&self) -> Option<&Arc<SymbolTable>> {
+        self.0.input_symbols()
+    }
+
+    fn output_symbols(&self) -> Option<&Arc<SymbolTable>> {
+        self.0.output_symbols()
+    }
+
+    fn set_input_symbols(&mut self, symt: Arc<SymbolTable>) {
+        self.0.set_input_symbols(symt)
+    }
+
+    fn set_output_symbols(&mut self, symt: Arc<SymbolTable>) {
+        self.0.set_output_symbols(symt)
+    }
+
+    fn take_input_symbols(&mut self) -> Option<Arc<SymbolTable>> {
+        self.0.take_input_symbols()
+    }
+
+    fn take_output_symbols(&mut self) -> Option<Arc<SymbolTable>> {
+        self.0.take_output_symbols()
+    }
+}
+
+impl<W, F, B> Debug for LookaheadRelabelFst<W, F, B>
+where
+    W: Semiring,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::fst_impls::VectorFst;
+    use crate::semirings::TropicalWeight;
+
+    use super::*;
+
+    #[test]
+    fn test_lookahead_relabel_fst_sync() {
+        fn is_sync<T: Sync>() {}
+        is_sync::<LookaheadRelabelFst<TropicalWeight, VectorFst<_>, VectorFst<_>>>();
+    }
+}
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
new file mode 100644
index 000000000..1862e1c93
--- /dev/null
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
@@ -0,0 +1,77 @@
+use std::borrow::Borrow;
+use std::marker::PhantomData;
+
+use anyhow::Result;
+
+use crate::algorithms::lazy::FstOp;
+use crate::algorithms::tr_compares::TrCompare;
+use crate::fst_properties::FstProperties;
+use crate::fst_traits::Fst;
+use crate::prelude::compose::LabelReachableData;
+use crate::prelude::{ILabelCompare, OLabelCompare};
+use crate::{Semiring, Trs, TrsVec};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+pub struct LookaheadRelabelFstOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
+    fst: B,
+    label_reachable_data: LabelReachableData,
+    relabel_input: bool,
+    ghost: PhantomData<(W, F)>,
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFstOp<W, F, B> {
+    pub fn new(fst: B, label_reachable_data: LabelReachableData, relabel_input: bool) -> Self {
+        Self {
+            fst,
+            label_reachable_data,
+            relabel_input,
+            ghost: PhantomData,
+        }
+    }
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> Debug for LookaheadRelabelFstOp<W, F, B> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "LookaheadRelabelFstOp {{ fst: {:?}, label_reachable_data: {:?}, relabel_input: {:?} }}", self.fst.borrow(), self.label_reachable_data, self.relabel_input)
+    }
+}
+
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> FstOp<W> for LookaheadRelabelFstOp<W, F, B> {
+    fn compute_start(&self) -> Result<Option<u32>> {
+        Ok(self.fst.borrow().start())
+    }
+
+    fn compute_trs(&self, id: u32) -> Result<TrsVec<W>> {
+        let trs = self.fst.borrow().get_trs(id)?;
+        let mut new_trs = Vec::with_capacity(trs.trs().len());
+        for tr in trs.trs() {
+            let mut new_tr = tr.clone();
+            if self.relabel_input {
+                new_tr.ilabel = self.label_reachable_data.relabel_unmut(tr.ilabel)?;
+            } else {
+                new_tr.olabel = self.label_reachable_data.relabel_unmut(tr.olabel)?;
+            }
+            new_trs.push(new_tr);
+        }
+
+        if self.relabel_input {
+            new_trs.sort_by(ILabelCompare::compare)
+        } else {
+            new_trs.sort_by(OLabelCompare::compare)
+        }
+        Ok(TrsVec(Arc::new(new_trs)))
+    }
+
+    fn compute_final_weight(&self, id: u32) -> Result<Option<W>> {
+        self.fst.borrow().final_weight(id)
+    }
+
+    fn properties(&self) -> FstProperties {
+        if self.relabel_input {
+            FstProperties::I_LABEL_SORTED
+        } else {
+            FstProperties::O_LABEL_SORTED
+        }
+    }
+}
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/mod.rs b/rustfst/src/algorithms/compose/lookahead_relabel/mod.rs
new file mode 100644
index 000000000..e8def26c1
--- /dev/null
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/mod.rs
@@ -0,0 +1,4 @@
+mod lookahead_relabel_fst;
+mod lookahead_relabel_fst_op;
+
+pub use self::lookahead_relabel_fst::LookaheadRelabelFst;
diff --git a/rustfst/src/algorithms/compose/matcher_fst.rs b/rustfst/src/algorithms/compose/matcher_fst.rs
index fcf9f7d5f..09c13f14c 100644
--- a/rustfst/src/algorithms/compose/matcher_fst.rs
+++ b/rustfst/src/algorithms/compose/matcher_fst.rs
@@ -93,29 +93,32 @@ where
         })
     }
 
-    // // Construct a new Matcher Fst intended for LookAhead composition and relabel fst2 wrt to the first fst.
-    // pub fn new_with_relabeling_2<F2: Fst<W>>(
-    //     mut fst: F,
-    //     fst2: &F2,
-    //     relabel_input: bool,
-    // ) -> Result<Self> {
-    //     let imatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchInput)?;
-    //     let omatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchOutput)?;
-    //
-    //     let mut add_on = (imatcher_data, omatcher_data);
-    //
-    //     LabelLookAheadRelabeler::init(&mut fst, &mut add_on)?;
-    //     LabelLookAheadRelabeler::relabel(fst2, &mut add_on, relabel_input)?;
-    //
-    //     let add_on = (add_on.0.map(Arc::new), add_on.1.map(Arc::new));
-    //
-    //     let fst_add_on = FstAddOn::new(fst, add_on);
-    //     Ok(Self {
-    //         fst_add_on,
-    //         matcher: PhantomData,
-    //         w: PhantomData,
-    //     })
-    // }
+    // Construct a new Matcher Fst intended for LookAhead composition and relabel fst2 wrt to the first fst.
+    pub fn new_with_relabeling_lazy_right<F2: Fst<W> + 'static>(
+        mut fst: F,
+        fst2: Arc<F2>,
+        relabel_input: bool,
+    ) -> Result<(Self, impl Fst<W>)> {
+        let imatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchInput)?;
+        let omatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchOutput)?;
+
+        let mut add_on = (imatcher_data, omatcher_data);
+
+        LabelLookAheadRelabeler::init(&mut fst, &mut add_on)?;
+        let la_fst2 = LabelLookAheadRelabeler::relabel_lazy(fst2, &mut add_on, relabel_input)?;
+
+        let add_on = (add_on.0.map(Arc::new), add_on.1.map(Arc::new));
+
+        let fst_add_on = FstAddOn::new(fst, add_on);
+        Ok((
+            Self {
+                fst_add_on,
+                matcher: PhantomData,
+                w: PhantomData,
+            },
+            la_fst2,
+        ))
+    }
 }
 
 impl<W: Semiring, F: CoreFst<W>, B: Borrow<F>, M, T> CoreFst<W> for MatcherFst<W, F, B, M, T> {
diff --git a/rustfst/src/algorithms/compose/mod.rs b/rustfst/src/algorithms/compose/mod.rs
index 627d2ccf6..a68672852 100644
--- a/rustfst/src/algorithms/compose/mod.rs
+++ b/rustfst/src/algorithms/compose/mod.rs
@@ -28,5 +28,8 @@ mod compose_static;
 mod interval_reach_visitor;
 mod interval_set;
 mod label_reachable;
+mod lookahead_relabel;
 mod matcher_fst;
 mod state_reachable;
+
+pub use self::lookahead_relabel::LookaheadRelabelFst;
diff --git a/rustfst/src/algorithms/lazy/lazy_fst.rs b/rustfst/src/algorithms/lazy/lazy_fst.rs
index 60b0778ea..b9ff14557 100644
--- a/rustfst/src/algorithms/lazy/lazy_fst.rs
+++ b/rustfst/src/algorithms/lazy/lazy_fst.rs
@@ -61,9 +61,11 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     fn num_trs(&self, s: StateId) -> Result<usize> {
-        self.cache
-            .num_trs(s)
-            .ok_or_else(|| format_err!("State {:?} doesn't exist", s))
+        let trs = self.get_trs(s)?;
+        Ok(trs.trs().len())
+        // self.cache
+        //     .num_trs(s)
+        //     .ok_or_else(|| format_err!("NT State {:?} doesn't exist", s))
     }
 
     unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
@@ -92,13 +94,13 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     fn num_input_epsilons(&self, state: StateId) -> Result<usize> {
         self.cache
             .num_input_epsilons(state)
-            .ok_or_else(|| format_err!("State {:?} doesn't exist", state))
+            .ok_or_else(|| format_err!("NI State {:?} doesn't exist", state))
     }
 
     fn num_output_epsilons(&self, state: StateId) -> Result<usize> {
         self.cache
             .num_output_epsilons(state)
-            .ok_or_else(|| format_err!("State {:?} doesn't exist", state))
+            .ok_or_else(|| format_err!("NO State {:?} doesn't exist", state))
     }
 }
 
diff --git a/rustfst/src/algorithms/relabel/mod.rs b/rustfst/src/algorithms/relabel/mod.rs
index 142ad1260..e34c31bc4 100644
--- a/rustfst/src/algorithms/relabel/mod.rs
+++ b/rustfst/src/algorithms/relabel/mod.rs
@@ -1,2 +1,4 @@
 mod relabel_fst;
 mod relabel_fst_op;
+
+pub use self::relabel_fst::RelabelFst;
diff --git a/rustfst/src/algorithms/relabel/relabel_fst.rs b/rustfst/src/algorithms/relabel/relabel_fst.rs
index 964aa9d7d..17c157d93 100644
--- a/rustfst/src/algorithms/relabel/relabel_fst.rs
+++ b/rustfst/src/algorithms/relabel/relabel_fst.rs
@@ -1,17 +1,17 @@
 use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache};
+use crate::algorithms::relabel::relabel_fst_op::RelabelFstOp;
+use crate::algorithms::relabel_pairs::iterator_to_hashmap;
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
 use crate::{Semiring, StateId, SymbolTable, TrsVec};
-use anyhow::{Result, Context};
+use anyhow::{Context, Result};
 use std::borrow::Borrow;
 use std::fmt::Debug;
 use std::sync::Arc;
-use crate::algorithms::relabel::relabel_fst_op::RelabelFstOp;
-use crate::algorithms::relabel_pairs::iterator_to_hashmap;
 
 type InnerLazyFst<W, F, B> = LazyFst<W, RelabelFstOp<W, F, B>, SimpleHashMapCache<W>>;
 
-struct RelabelFst<W: Semiring, F: Fst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
+pub struct RelabelFst<W: Semiring, F: Fst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
 
 impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RelabelFst<W, F, B> {
     pub fn new<I, J>(fst: B, ipairs: I, opairs: J) -> Result<Self>
diff --git a/rustfst/src/algorithms/relabel/relabel_fst_op.rs b/rustfst/src/algorithms/relabel/relabel_fst_op.rs
index fc1931efb..8215b0a0c 100644
--- a/rustfst/src/algorithms/relabel/relabel_fst_op.rs
+++ b/rustfst/src/algorithms/relabel/relabel_fst_op.rs
@@ -6,9 +6,9 @@ use anyhow::Result;
 use crate::algorithms::lazy::FstOp;
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::Fst;
-use crate::{Semiring, TrsVec, StateId, Trs};
-use std::fmt::{Debug, Formatter};
+use crate::{Semiring, StateId, Trs, TrsVec};
 use std::collections::HashMap;
+use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
 pub struct RelabelFstOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
@@ -19,16 +19,29 @@ pub struct RelabelFstOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
 }
 
 impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RelabelFstOp<W, F, B> {
-    pub fn new(fst: B, map_ilabels: HashMap<StateId, StateId>, map_olabels: HashMap<StateId, StateId>) -> Self {
+    pub fn new(
+        fst: B,
+        map_ilabels: HashMap<StateId, StateId>,
+        map_olabels: HashMap<StateId, StateId>,
+    ) -> Self {
         Self {
-            fst, map_ilabels, map_olabels, ghost: PhantomData
+            fst,
+            map_ilabels,
+            map_olabels,
+            ghost: PhantomData,
         }
     }
 }
 
 impl<W: Semiring, F: Fst<W>, B: Borrow<F>> Debug for RelabelFstOp<W, F, B> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        unimplemented!()
+        write!(
+            f,
+            "RelabelFstOp {{ fst: {:?}, map_ilabels: {:?}, map_olabels: {:?} }}",
+            self.fst.borrow(),
+            self.map_ilabels,
+            self.map_olabels
+        )
     }
 }
 
@@ -48,6 +61,7 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> FstOp<W> for RelabelFstOp<W, F, B> {
             if let Some(new_olabel) = self.map_olabels.get(&tr.olabel) {
                 new_tr.olabel = *new_olabel;
             }
+            trs.push(new_tr);
         }
         Ok(TrsVec(Arc::new(trs)))
     }
diff --git a/rustfst/src/tests_openfst/algorithms/compose.rs b/rustfst/src/tests_openfst/algorithms/compose.rs
index 7615ab309..ae5e9cba8 100644
--- a/rustfst/src/tests_openfst/algorithms/compose.rs
+++ b/rustfst/src/tests_openfst/algorithms/compose.rs
@@ -23,10 +23,12 @@ use crate::algorithms::compose::{ComposeFilterEnum, ComposeFst, ComposeFstOpOpti
 use crate::algorithms::lazy::SimpleHashMapCache;
 use crate::algorithms::{tr_compares::ILabelCompare, tr_sort};
 use crate::fst_impls::VectorFst;
-use crate::fst_traits::SerializableFst;
+use crate::fst_traits::{CoreFst, ExpandedFst, Fst, MutableFst, SerializableFst};
 use crate::semirings::{SerializableSemiring, WeaklyDivisibleSemiring, WeightQuantize};
 use crate::tests_openfst::utils::test_eq_fst;
 use crate::tests_openfst::FstTestData;
+use crate::{StateId, SymbolTable, Trs};
+use std::cmp::max;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct ComposeOperationResult {
@@ -253,6 +255,156 @@ where
     Ok(())
 }
 
+fn do_test_compose_lookahead_lazy_right<W>(
+    fst_raw: &VectorFst<W>,
+    compose_test_data: &ComposeTestData<W, VectorFst<W>>,
+) -> Result<()>
+where
+    W: SerializableSemiring + WeightQuantize + WeaklyDivisibleSemiring,
+{
+    println!("Lookahead composition lazy right");
+    type TLaFst<S, F> = MatcherFst<
+        S,
+        F,
+        Arc<F>,
+        LabelLookAheadMatcher<
+            S,
+            F,
+            Arc<F>,
+            SortedMatcher<S, F, Arc<F>>,
+            DefaultLabelLookAheadMatcherFlags,
+        >,
+        LabelReachableData,
+    >;
+
+    type TMatcher1<S, F> = LabelLookAheadMatcher<
+        S,
+        F,
+        Arc<F>,
+        SortedMatcher<S, F, Arc<F>>,
+        DefaultLabelLookAheadMatcherFlags,
+    >;
+    type TMatcher2<S, F> = SortedMatcher<S, F, Arc<F>>;
+
+    type TSeqFilter<S, F1, F2> = AltSequenceComposeFilterBuilder<
+        S,
+        F1,
+        F2,
+        Arc<F1>,
+        Arc<F2>,
+        TMatcher1<S, F1>,
+        TMatcher2<S, F2>,
+    >;
+    type TLookFilter<S, F1, F2> = LookAheadComposeFilterBuilder<
+        S,
+        F1,
+        F2,
+        Arc<F1>,
+        Arc<F2>,
+        TMatcher1<S, F1>,
+        TMatcher2<S, F2>,
+        TSeqFilter<S, F1, F2>,
+        SMatchOutput,
+    >;
+    type TPushWeightsFilter<S, F1, F2> = PushWeightsComposeFilterBuilder<
+        S,
+        F1,
+        F2,
+        Arc<F1>,
+        Arc<F2>,
+        TMatcher1<S, F1>,
+        TMatcher2<S, F2>,
+        TLookFilter<S, F1, F2>,
+        SMatchOutput,
+    >;
+    type TPushLabelsFilter<S, F1, F2> = PushLabelsComposeFilterBuilder<
+        S,
+        F1,
+        F2,
+        Arc<F1>,
+        Arc<F2>,
+        TMatcher1<S, F1>,
+        TMatcher2<S, F2>,
+        TPushWeightsFilter<S, F1, F2>,
+        SMatchOutput,
+    >;
+
+    type TComposeFilter<S, F1, F2> = TPushLabelsFilter<S, F1, F2>;
+
+    let fst1: VectorFst<_> = fst_raw.clone().into();
+    let mut fst2: VectorFst<_> = compose_test_data.fst_2.clone();
+
+    let mut symt = SymbolTable::new();
+    let mut max_ilabel = 0;
+    for s in 0..fst2.num_states() {
+        let trs = fst2.get_trs(s as StateId)?;
+        for tr in trs.trs() {
+            max_ilabel = max(max_ilabel, tr.ilabel);
+        }
+    }
+    for i in 1..=max_ilabel {
+        symt.add_symbol(i.to_string());
+    }
+    fst2.set_input_symbols(Arc::new(symt));
+
+    let (graph1look, lazy_fst2) =
+        TLaFst::new_with_relabeling_lazy_right(fst1, Arc::new(fst2), true)?;
+    let graph1look = Arc::new(graph1look);
+    let fst2 = Arc::new(lazy_fst2);
+
+    // LabelLookAheadRelabeler::relabel(&mut fst2, graph1look.addon(), true)?;
+
+    // tr_sort(&mut fst2, ILabelCompare {});
+
+    // let fst2 = Arc::new(fst2);
+
+    let matcher1 = TMatcher1::new_with_data(
+        Arc::clone(&graph1look),
+        MatchType::MatchOutput,
+        graph1look.data(MatchType::MatchOutput).cloned(),
+    )?;
+
+    let matcher2 = TMatcher2::new(Arc::clone(&fst2), MatchType::MatchInput)?;
+
+    let compose_filter = TComposeFilter::new(
+        Arc::clone(&graph1look),
+        Arc::clone(&fst2),
+        Some(matcher1),
+        Some(matcher2),
+    )?;
+
+    let compose_options = ComposeFstOpOptions::<_, _, TComposeFilter<_, _, _>, _>::new(
+        // compose_filter.matcher1(),
+        None,
+        // compose_filter.matcher2(),
+        None,
+        compose_filter,
+        None,
+    );
+
+    let dyn_fst = ComposeFst::<_, _, _, _, _, _, _, _, SimpleHashMapCache<_>>::new_with_options(
+        graph1look,
+        fst2,
+        compose_options,
+    )?;
+
+    // Check clonability
+    fn is_clone<T: Clone>(_v: &T) {}
+    is_clone(&dyn_fst);
+
+    let mut static_fst: VectorFst<_> = dyn_fst.compute()?;
+    // TODO: This is a hotfix that should be removed.
+    static_fst.compute_and_update_properties_all()?;
+
+    test_eq_fst(
+        &compose_test_data.result,
+        &static_fst,
+        format!("Compose failed : filter_name = lookahead_lazy_right"),
+    );
+
+    Ok(())
+}
+
 pub fn test_compose<W>(test_data: &FstTestData<W, VectorFst<W>>) -> Result<()>
 where
     W: SerializableSemiring + WeightQuantize + WeaklyDivisibleSemiring,
@@ -294,7 +446,10 @@ where
                 compose_test_data,
                 ComposeFilterEnum::NoMatchFilter,
             )?,
-            "lookahead" => do_test_compose_lookahead(&test_data.raw, compose_test_data)?,
+            "lookahead" => {
+                do_test_compose_lookahead(&test_data.raw, compose_test_data)?;
+                do_test_compose_lookahead_lazy_right(&test_data.raw, compose_test_data)?;
+            }
             _ => panic!("Not supported : {}", &compose_test_data.filter_name),
         }
     }

From 650c270b48fda7fbdeb115e1f3982b0360049041 Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 20 Dec 2022 11:27:16 +0100
Subject: [PATCH 3/9] Hot fix

---
 .../lookahead_relabel_fst.rs                  |  5 ++-
 .../lookahead_relabel_fst_op.rs               |  2 +-
 rustfst/src/algorithms/compose/matcher_fst.rs |  2 -
 rustfst/src/algorithms/lazy/lazy_fst.rs       | 38 ++++++++++++++++---
 rustfst/src/algorithms/lazy/mod.rs            |  2 +-
 rustfst/src/fst_traits/expanded_fst.rs        |  3 ++
 .../src/tests_openfst/algorithms/compose.rs   |  4 ++
 7 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
index a4006c7b6..346b889bc 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -1,5 +1,5 @@
 use crate::algorithms::compose::lookahead_relabel::lookahead_relabel_fst_op::LookaheadRelabelFstOp;
-use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache};
+use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache, iterate_lazy};
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
 use crate::prelude::compose::LabelReachableData;
@@ -35,6 +35,9 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFst<W, F, B> {
     }
 
     pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
+        // let underneath_fst = self.0.op.fst.borrow();
+        // Trick to compute the underneath lazy fst
+        // iterate_lazy(underneath_fst)?;
         self.0.compute()
     }
 }
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
index 1862e1c93..572c81735 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst_op.rs
@@ -14,7 +14,7 @@ use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
 pub struct LookaheadRelabelFstOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
-    fst: B,
+    pub(crate) fst: B,
     label_reachable_data: LabelReachableData,
     relabel_input: bool,
     ghost: PhantomData<(W, F)>,
diff --git a/rustfst/src/algorithms/compose/matcher_fst.rs b/rustfst/src/algorithms/compose/matcher_fst.rs
index 09c13f14c..b1fa147ee 100644
--- a/rustfst/src/algorithms/compose/matcher_fst.rs
+++ b/rustfst/src/algorithms/compose/matcher_fst.rs
@@ -103,10 +103,8 @@ where
         let omatcher_data = M::create_data::<F, _>(&fst, MatchType::MatchOutput)?;
 
         let mut add_on = (imatcher_data, omatcher_data);
-
         LabelLookAheadRelabeler::init(&mut fst, &mut add_on)?;
         let la_fst2 = LabelLookAheadRelabeler::relabel_lazy(fst2, &mut add_on, relabel_input)?;
-
         let add_on = (add_on.0.map(Arc::new), add_on.1.map(Arc::new));
 
         let fst_add_on = FstAddOn::new(fst, add_on);
diff --git a/rustfst/src/algorithms/lazy/lazy_fst.rs b/rustfst/src/algorithms/lazy/lazy_fst.rs
index b9ff14557..ed1ea3835 100644
--- a/rustfst/src/algorithms/lazy/lazy_fst.rs
+++ b/rustfst/src/algorithms/lazy/lazy_fst.rs
@@ -61,11 +61,11 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     fn num_trs(&self, s: StateId) -> Result<usize> {
-        let trs = self.get_trs(s)?;
-        Ok(trs.trs().len())
-        // self.cache
-        //     .num_trs(s)
-        //     .ok_or_else(|| format_err!("NT State {:?} doesn't exist", s))
+        // let trs = self.get_trs(s)?;
+        // Ok(trs.trs().len())
+        self.cache
+            .num_trs(s)
+            .ok_or_else(|| format_err!("NT State {:?} doesn't exist", s))
     }
 
     unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
@@ -272,6 +272,34 @@ where
     }
 }
 
+/// Turns the Lazy FST into a static one.
+pub fn iterate_lazy<W: Semiring, F: Fst<W>>(fst: &F) -> Result<()> {
+    let start_state = fst.start();
+    let start_state = match start_state {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+    let mut queue = VecDeque::new();
+    let mut visited_states = vec![];
+    visited_states.resize(start_state as usize + 1, false);
+    visited_states[start_state as usize] = true;
+    queue.push_back(start_state);
+    while let Some(s) = queue.pop_front() {
+        let trs_owner = fst.get_trs(s)?;
+        for tr in trs_owner.trs() {
+            if (tr.nextstate as usize) >= visited_states.len() {
+                visited_states.resize(tr.nextstate as usize + 1, false);
+            }
+            if !visited_states[tr.nextstate as usize] {
+                queue.push_back(tr.nextstate);
+                visited_states[tr.nextstate as usize] = true;
+            }
+        }
+        fst.final_weight(s)?;
+    }
+    Ok(())
+}
+
 impl<W, Op, Cache> SerializableLazyFst for LazyFst<W, Op, Cache>
 where
     W: SerializableSemiring,
diff --git a/rustfst/src/algorithms/lazy/mod.rs b/rustfst/src/algorithms/lazy/mod.rs
index 68368fab8..cd56766ed 100644
--- a/rustfst/src/algorithms/lazy/mod.rs
+++ b/rustfst/src/algorithms/lazy/mod.rs
@@ -1,7 +1,7 @@
 pub use self::cache::*;
 pub use fst_op::{AccessibleOpState, FstOp, SerializableOpState};
 pub use fst_op_2::FstOp2;
-pub use lazy_fst::{LazyFst, SerializableLazyFst};
+pub use lazy_fst::{LazyFst, SerializableLazyFst, iterate_lazy};
 pub use lazy_fst_2::LazyFst2;
 pub use state_table::StateTable;
 
diff --git a/rustfst/src/fst_traits/expanded_fst.rs b/rustfst/src/fst_traits/expanded_fst.rs
index 66f4c1108..17720bfc3 100644
--- a/rustfst/src/fst_traits/expanded_fst.rs
+++ b/rustfst/src/fst_traits/expanded_fst.rs
@@ -60,10 +60,12 @@ pub trait ExpandedFst<W: Semiring>: Fst<W> + Clone + PartialEq + FstIntoIterator
                     || tr1.olabel != tr2.olabel
                     || tr1.nextstate != tr2.nextstate
                 {
+                    println!("State {} Ilabel {} {} Olabel {} {} Nextstate {} {}", state, tr1.ilabel, tr2.ilabel, tr1.olabel, tr2.olabel, tr1.nextstate, tr2.nextstate);
                     return false;
                 }
 
                 if !tr1.weight.approx_equal(&tr2.weight, delta) {
+                    println!("Not the same weight : {:?} != {:?}", tr1.weight, tr2.weight);
                     return false;
                 }
             }
@@ -79,6 +81,7 @@ pub trait ExpandedFst<W: Semiring>: Fst<W> + Clone + PartialEq + FstIntoIterator
             };
 
             if !fw_equal {
+                println!("Pouet");
                 return false;
             }
         }
diff --git a/rustfst/src/tests_openfst/algorithms/compose.rs b/rustfst/src/tests_openfst/algorithms/compose.rs
index ae5e9cba8..34e8f1310 100644
--- a/rustfst/src/tests_openfst/algorithms/compose.rs
+++ b/rustfst/src/tests_openfst/algorithms/compose.rs
@@ -29,6 +29,7 @@ use crate::tests_openfst::utils::test_eq_fst;
 use crate::tests_openfst::FstTestData;
 use crate::{StateId, SymbolTable, Trs};
 use std::cmp::max;
+use crate::prelude::lazy::SerializableLazyFst;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct ComposeOperationResult {
@@ -392,6 +393,9 @@ where
     fn is_clone<T: Clone>(_v: &T) {}
     is_clone(&dyn_fst);
 
+    fn is_serializable<T: SerializableLazyFst>(_v: &T) {}
+    is_serializable(&dyn_fst);
+
     let mut static_fst: VectorFst<_> = dyn_fst.compute()?;
     // TODO: This is a hotfix that should be removed.
     static_fst.compute_and_update_properties_all()?;

From 5a18530693f18d45e42d81e62d1ad91d9ad0af0d Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 20 Dec 2022 17:37:21 +0100
Subject: [PATCH 4/9] Fix lazy compute

---
 rustfst-ffi/src/algorithms/compose.rs         |  4 ++--
 rustfst-ffi/src/fst/mod.rs                    |  2 +-
 rustfst-ffi/src/iterators.rs                  |  6 +++---
 rustfst-ffi/src/string_path.rs                |  2 +-
 rustfst-ffi/src/string_paths_iterator.rs      |  2 +-
 rustfst-ffi/src/symbol_table.rs               |  2 +-
 rustfst-ffi/src/tr.rs                         |  2 +-
 rustfst-ffi/src/trs.rs                        |  2 +-
 rustfst/src/algorithms/compose/compose_fst.rs | 12 ++++++++---
 .../src/algorithms/compose/compose_fst_op.rs  |  4 ++--
 .../compose/interval_reach_visitor.rs         | 13 +++++++-----
 .../lookahead_relabel_fst.rs                  |  4 ++--
 .../src/algorithms/compose/state_reachable.rs | 12 +++++------
 rustfst/src/algorithms/condense.rs            |  2 +-
 rustfst/src/algorithms/connect.rs             |  6 +++---
 rustfst/src/algorithms/dfs_visit.rs           | 21 +++++++++++--------
 rustfst/src/algorithms/lazy/mod.rs            |  2 +-
 rustfst/src/algorithms/queues/auto_queue.rs   |  4 ++--
 .../src/algorithms/queues/top_order_queue.rs  | 12 +++++++----
 rustfst/src/algorithms/randgen/mod.rs         |  2 +-
 .../src/algorithms/randgen/randgen_visitor.rs |  4 ++--
 .../rm_epsilon/rm_epsilon_static.rs           |  4 ++--
 rustfst/src/algorithms/rm_final_epsilon.rs    |  2 +-
 rustfst/src/algorithms/top_sort.rs            |  6 +++---
 .../src/algorithms/visitors/scc_visitors.rs   |  5 +++--
 .../fst_properties/compute_fst_properties.rs  |  2 +-
 rustfst/src/fst_traits/expanded_fst.rs        | 11 +++++++++-
 .../src/tests_openfst/algorithms/compose.rs   |  2 +-
 28 files changed, 89 insertions(+), 63 deletions(-)

diff --git a/rustfst-ffi/src/algorithms/compose.rs b/rustfst-ffi/src/algorithms/compose.rs
index 56782d39c..f0cd7717b 100644
--- a/rustfst-ffi/src/algorithms/compose.rs
+++ b/rustfst-ffi/src/algorithms/compose.rs
@@ -269,7 +269,7 @@ pub extern "C" fn fst_matcher_config_destroy(ptr: *mut CMatcherConfig) -> RUSTFS
         }
 
         unsafe {
-            Box::from_raw(ptr);
+            drop(Box::from_raw(ptr));
         }
         Ok(())
     })
@@ -283,7 +283,7 @@ pub extern "C" fn fst_compose_config_destroy(ptr: *mut CComposeConfig) -> RUSTFS
         }
 
         unsafe {
-            Box::from_raw(ptr);
+            drop(Box::from_raw(ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/fst/mod.rs b/rustfst-ffi/src/fst/mod.rs
index 4eef0e4fe..548b61215 100644
--- a/rustfst-ffi/src/fst/mod.rs
+++ b/rustfst-ffi/src/fst/mod.rs
@@ -327,7 +327,7 @@ pub fn fst_destroy(fst_ptr: *mut CFst) -> RUSTFST_FFI_RESULT {
         }
 
         unsafe {
-            Box::from_raw(fst_ptr);
+            drop(Box::from_raw(fst_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/iterators.rs b/rustfst-ffi/src/iterators.rs
index 4d2756fb3..045e667fa 100644
--- a/rustfst-ffi/src/iterators.rs
+++ b/rustfst-ffi/src/iterators.rs
@@ -113,7 +113,7 @@ pub extern "C" fn trs_iterator_destroy(iter_ptr: *mut CTrsIterator) -> RUSTFST_F
         }
 
         unsafe {
-            Box::from_raw(iter_ptr);
+            drop(Box::from_raw(iter_ptr));
         }
         Ok(())
     })
@@ -274,7 +274,7 @@ pub extern "C" fn mut_trs_iterator_destroy(iter_ptr: *mut CMutTrsIterator) -> RU
         }
 
         unsafe {
-            Box::from_raw(iter_ptr);
+            drop(Box::from_raw(iter_ptr));
         }
         Ok(())
     })
@@ -333,7 +333,7 @@ pub extern "C" fn state_iterator_destroy(iter_ptr: *mut CStateIterator) -> RUSTF
         }
 
         unsafe {
-            Box::from_raw(iter_ptr);
+            drop(Box::from_raw(iter_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/string_path.rs b/rustfst-ffi/src/string_path.rs
index 0fcf6339b..fec6a5236 100644
--- a/rustfst-ffi/src/string_path.rs
+++ b/rustfst-ffi/src/string_path.rs
@@ -17,7 +17,7 @@ pub extern "C" fn string_path_destroy(iter_ptr: *mut CStringPath) -> RUSTFST_FFI
         }
 
         unsafe {
-            Box::from_raw(iter_ptr);
+            drop(Box::from_raw(iter_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/string_paths_iterator.rs b/rustfst-ffi/src/string_paths_iterator.rs
index 4027ccfd7..1d6897298 100644
--- a/rustfst-ffi/src/string_paths_iterator.rs
+++ b/rustfst-ffi/src/string_paths_iterator.rs
@@ -105,7 +105,7 @@ pub extern "C" fn string_paths_iterator_destroy(
         }
 
         unsafe {
-            Box::from_raw(iter_ptr);
+            drop(Box::from_raw(iter_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/symbol_table.rs b/rustfst-ffi/src/symbol_table.rs
index aaedec548..5dfb240e5 100644
--- a/rustfst-ffi/src/symbol_table.rs
+++ b/rustfst-ffi/src/symbol_table.rs
@@ -206,7 +206,7 @@ pub extern "C" fn symt_destroy(symt_ptr: *mut CSymbolTable) -> RUSTFST_FFI_RESUL
         }
 
         unsafe {
-            Box::from_raw(symt_ptr);
+            drop(Box::from_raw(symt_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/tr.rs b/rustfst-ffi/src/tr.rs
index 315d955a8..5d0aa1048 100644
--- a/rustfst-ffi/src/tr.rs
+++ b/rustfst-ffi/src/tr.rs
@@ -150,7 +150,7 @@ pub extern "C" fn tr_delete(tr_ptr: *mut CTr) -> RUSTFST_FFI_RESULT {
         }
 
         unsafe {
-            Box::from_raw(tr_ptr);
+            drop(Box::from_raw(tr_ptr));
         }
         Ok(())
     })
diff --git a/rustfst-ffi/src/trs.rs b/rustfst-ffi/src/trs.rs
index bd7f647a2..a3c1c2feb 100644
--- a/rustfst-ffi/src/trs.rs
+++ b/rustfst-ffi/src/trs.rs
@@ -92,7 +92,7 @@ pub extern "C" fn trs_vec_delete(trs_ptr: *mut CTrs) -> RUSTFST_FFI_RESULT {
         }
 
         unsafe {
-            Box::from_raw(trs_ptr);
+            drop(Box::from_raw(trs_ptr));
         }
         Ok(())
     })
diff --git a/rustfst/src/algorithms/compose/compose_fst.rs b/rustfst/src/algorithms/compose/compose_fst.rs
index 493955a3b..72c411ba6 100644
--- a/rustfst/src/algorithms/compose/compose_fst.rs
+++ b/rustfst/src/algorithms/compose/compose_fst.rs
@@ -1,7 +1,9 @@
-use anyhow::Result;
 use std::borrow::Borrow;
 use std::fmt::Debug;
 use std::path::Path;
+use std::sync::Arc;
+
+use anyhow::Result;
 
 use crate::algorithms::compose::compose_filters::{
     ComposeFilter, ComposeFilterBuilder, SequenceComposeFilterBuilder,
@@ -10,6 +12,7 @@ use crate::algorithms::compose::matchers::{GenericMatcher, Matcher};
 use crate::algorithms::compose::{
     ComposeFstOp, ComposeFstOpOptions, ComposeFstOpState, ComposeStateTuple,
 };
+use crate::algorithms::lazy::iterate_lazy;
 use crate::algorithms::lazy::{
     FstCache, LazyFst, SerializableCache, SerializableLazyFst, SimpleVecCache,
 };
@@ -18,7 +21,6 @@ use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, Sta
 use crate::parsers::SerializeBinary;
 use crate::semirings::{Semiring, SerializableSemiring};
 use crate::{StateId, SymbolTable, TrsVec};
-use std::sync::Arc;
 
 type InnerLazyFst<W, F1, F2, B1, B2, M1, M2, CFB, Cache> =
     LazyFst<W, ComposeFstOp<W, F1, F2, B1, B2, M1, M2, CFB>, Cache>;
@@ -160,6 +162,9 @@ where
 
     /// Turns the Lazy FST into a static one.
     pub fn compute<F: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F> {
+        // Small trick to make sure that both FSTs are fully expanded.
+        iterate_lazy(self.0.op.fst1.borrow())?;
+        iterate_lazy(self.0.op.fst2.borrow())?;
         self.0.compute()
     }
 }
@@ -356,11 +361,12 @@ where
 
 #[cfg(test)]
 mod test {
-    use super::*;
     use crate::algorithms::compose::matchers::SortedMatcher;
     use crate::fst_impls::VectorFst;
     use crate::semirings::TropicalWeight;
 
+    use super::*;
+
     #[test]
     fn test_compose_fst_sync() {
         fn is_sync<T: Sync>() {}
diff --git a/rustfst/src/algorithms/compose/compose_fst_op.rs b/rustfst/src/algorithms/compose/compose_fst_op.rs
index 28f1dfb8d..bbbd3f08d 100644
--- a/rustfst/src/algorithms/compose/compose_fst_op.rs
+++ b/rustfst/src/algorithms/compose/compose_fst_op.rs
@@ -82,8 +82,8 @@ where
     >,
     match_type: MatchType,
     properties: FstProperties,
-    fst1: B1,
-    fst2: B2,
+    pub(crate) fst1: B1,
+    pub(crate) fst2: B2,
 }
 
 impl<W, F1, F2, B1, B2, M1, M2, CFB> Clone for ComposeFstOp<W, F1, F2, B1, B2, M1, M2, CFB>
diff --git a/rustfst/src/algorithms/compose/interval_reach_visitor.rs b/rustfst/src/algorithms/compose/interval_reach_visitor.rs
index 0a24d42bf..f56f21001 100644
--- a/rustfst/src/algorithms/compose/interval_reach_visitor.rs
+++ b/rustfst/src/algorithms/compose/interval_reach_visitor.rs
@@ -1,9 +1,12 @@
+use std::cmp::Ordering;
+
+use anyhow::Result;
+
 use crate::algorithms::compose::{IntInterval, IntervalSet};
 use crate::algorithms::dfs_visit::Visitor;
 use crate::fst_traits::Fst;
 use crate::semirings::Semiring;
 use crate::{StateId, Tr};
-use std::cmp::Ordering;
 
 static UNASSIGNED: usize = std::usize::MAX;
 
@@ -30,18 +33,18 @@ impl<'a, W: Semiring, F: Fst<W>> Visitor<'a, W, F> for IntervalReachVisitor<'a,
     fn init_visit(&mut self, _fst: &'a F) {}
 
     /// Invoked when state discovered (2nd arg is DFS tree root).
-    fn init_state(&mut self, s: StateId, _root: StateId) -> bool {
+    fn init_state(&mut self, s: StateId, _root: StateId) -> Result<bool> {
         while self.isets.len() <= (s as usize) {
             self.isets.push(IntervalSet::default());
         }
         while self.state2index.len() <= (s as usize) {
             self.state2index.push(UNASSIGNED);
         }
-        if let Some(final_weight) = self.fst.final_weight(s).unwrap() {
+        if let Some(final_weight) = self.fst.final_weight(s)? {
             if !final_weight.is_zero() {
                 let interval_set = &mut self.isets[s as usize];
                 if self.index == UNASSIGNED {
-                    if self.fst.num_trs(s).unwrap() > 0 {
+                    if self.fst.num_trs(s)? > 0 {
                         panic!("IntervalReachVisitor: state2index map must be empty for this FST")
                     }
                     let index = self.state2index[s as usize];
@@ -56,7 +59,7 @@ impl<'a, W: Semiring, F: Fst<W>> Visitor<'a, W, F> for IntervalReachVisitor<'a,
                 }
             }
         }
-        true
+        Ok(true)
     }
 
     /// Invoked when tree transition to white/undiscovered state examined.
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
index 346b889bc..a574f2cfa 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -35,9 +35,9 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFst<W, F, B> {
     }
 
     pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
-        // let underneath_fst = self.0.op.fst.borrow();
+        let underneath_fst = self.0.op.fst.borrow();
         // Trick to compute the underneath lazy fst
-        // iterate_lazy(underneath_fst)?;
+        iterate_lazy(underneath_fst)?;
         self.0.compute()
     }
 }
diff --git a/rustfst/src/algorithms/compose/state_reachable.rs b/rustfst/src/algorithms/compose/state_reachable.rs
index c2cfffc43..5f7a17f32 100644
--- a/rustfst/src/algorithms/compose/state_reachable.rs
+++ b/rustfst/src/algorithms/compose/state_reachable.rs
@@ -27,7 +27,7 @@ impl StateReachable {
         let props = fst.properties_check(FstProperties::ACYCLIC)?;
         let acyclic = props.contains(FstProperties::ACYCLIC);
         if acyclic {
-            Ok(Self::new_acyclic(fst))
+            Ok(Self::new_acyclic(fst)?)
         } else {
             Self::new_cyclic(fst)
         }
@@ -35,7 +35,7 @@ impl StateReachable {
 
     pub fn new_cyclic<W: Semiring, F: ExpandedFst<W>>(fst: &F) -> Result<Self> {
         let (scc, cfst): (_, VectorFst<_>) = condense(fst).unwrap();
-        let reachable = StateReachable::new_acyclic(&cfst);
+        let reachable = StateReachable::new_acyclic(&cfst)?;
         let mut nscc = vec![];
 
         // Gets the number of states per SCC.
@@ -66,13 +66,13 @@ impl StateReachable {
         Ok(Self { isets, state2index })
     }
 
-    pub fn new_acyclic<W: Semiring, F: ExpandedFst<W>>(fst: &F) -> Self {
+    pub fn new_acyclic<W: Semiring, F: ExpandedFst<W>>(fst: &F) -> Result<Self> {
         let mut reach_visitor = IntervalReachVisitor::new(fst);
-        dfs_visit(fst, &mut reach_visitor, &AnyTrFilter {}, false);
-        Self {
+        dfs_visit(fst, &mut reach_visitor, &AnyTrFilter {}, false)?;
+        Ok(Self {
             isets: reach_visitor.isets,
             state2index: reach_visitor.state2index,
-        }
+        })
     }
 
     #[allow(unused)]
diff --git a/rustfst/src/algorithms/condense.rs b/rustfst/src/algorithms/condense.rs
index 18105f1bf..1ac2a4243 100644
--- a/rustfst/src/algorithms/condense.rs
+++ b/rustfst/src/algorithms/condense.rs
@@ -14,7 +14,7 @@ pub fn condense<W: Semiring, FI: Fst<W> + ExpandedFst<W>, FO: MutableFst<W>>(
     ifst: &FI,
 ) -> Result<(Vec<i32>, FO)> {
     let mut visitor = SccVisitor::new(ifst, true, false);
-    dfs_visit(ifst, &mut visitor, &AnyTrFilter {}, false);
+    dfs_visit(ifst, &mut visitor, &AnyTrFilter {}, false)?;
     let scc = visitor.scc.unwrap();
     let mut ofst = FO::new();
     if let Some(max) = scc.iter().max() {
diff --git a/rustfst/src/algorithms/connect.rs b/rustfst/src/algorithms/connect.rs
index 76a680f8a..2a476d13c 100644
--- a/rustfst/src/algorithms/connect.rs
+++ b/rustfst/src/algorithms/connect.rs
@@ -51,7 +51,7 @@ use crate::NO_STATE_ID;
 ///
 pub fn connect<W: Semiring, F: ExpandedFst<W> + MutableFst<W>>(fst: &mut F) -> Result<()> {
     let mut visitor = ConnectVisitor::new(fst);
-    dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false);
+    dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false)?;
     let mut dstates = Vec::with_capacity(visitor.access.len());
     for s in 0..visitor.access.len() {
         if !visitor.access[s] || !visitor.coaccess[s] {
@@ -100,7 +100,7 @@ impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> ConnectVisitor<'a, W, F> {
 impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> Visitor<'a, W, F> for ConnectVisitor<'a, W, F> {
     fn init_visit(&mut self, _fst: &'a F) {}
 
-    fn init_state(&mut self, s: StateId, root: StateId) -> bool {
+    fn init_state(&mut self, s: StateId, root: StateId) -> Result<bool> {
         self.scc_stack.push(s);
         let s = s as usize;
         self.dfnumber[s] = self.nstates as i32;
@@ -108,7 +108,7 @@ impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> Visitor<'a, W, F> for ConnectVisit
         self.onstack[s] = true;
         self.access[s] = root == self.start;
         self.nstates += 1;
-        true
+        Ok(true)
     }
 
     fn tree_tr(&mut self, _s: StateId, _tr: &Tr<W>) -> bool {
diff --git a/rustfst/src/algorithms/dfs_visit.rs b/rustfst/src/algorithms/dfs_visit.rs
index d810b2197..1cf3b69fa 100644
--- a/rustfst/src/algorithms/dfs_visit.rs
+++ b/rustfst/src/algorithms/dfs_visit.rs
@@ -1,12 +1,14 @@
+use std::marker::PhantomData;
+
+use anyhow::Result;
+use unsafe_unwrap::UnsafeUnwrap;
+
+use crate::algorithms::tr_filters::TrFilter;
 use crate::fst_traits::{ExpandedFst, Fst};
 use crate::semirings::Semiring;
 use crate::tr::Tr;
 use crate::{StateId, Trs};
 
-use crate::algorithms::tr_filters::TrFilter;
-use std::marker::PhantomData;
-use unsafe_unwrap::UnsafeUnwrap;
-
 #[derive(PartialOrd, PartialEq, Copy, Clone)]
 enum DfsStateColor {
     /// Undiscovered.
@@ -22,7 +24,7 @@ pub trait Visitor<'a, W: Semiring, F: Fst<W>> {
     fn init_visit(&mut self, fst: &'a F);
 
     /// Invoked when state discovered (2nd arg is DFS tree root).
-    fn init_state(&mut self, s: StateId, root: StateId) -> bool;
+    fn init_state(&mut self, s: StateId, root: StateId) -> Result<bool>;
 
     /// Invoked when tree transition to white/undiscovered state examined.
     fn tree_tr(&mut self, s: StateId, tr: &Tr<W>) -> bool;
@@ -100,12 +102,12 @@ pub fn dfs_visit<'a, W: Semiring, F: ExpandedFst<W>, V: Visitor<'a, W, F>, A: Tr
     visitor: &mut V,
     tr_filter: &A,
     access_only: bool,
-) {
+) -> Result<()> {
     visitor.init_visit(fst);
     let start = match fst.start() {
         None => {
             visitor.finish_visit();
-            return;
+            return Ok(());
         }
         Some(s) => s,
     };
@@ -123,7 +125,7 @@ pub fn dfs_visit<'a, W: Semiring, F: ExpandedFst<W>, V: Visitor<'a, W, F>, A: Tr
         }
         state_color[root as usize] = DfsStateColor::Grey;
         state_stack.push(DfsState::new(fst, root));
-        dfs = visitor.init_state(root, root);
+        dfs = visitor.init_state(root, root)?;
         let mut state_stack_next = None;
         while !state_stack.is_empty() {
             let dfs_state = unsafe { state_stack.last_mut().unsafe_unwrap() };
@@ -156,7 +158,7 @@ pub fn dfs_visit<'a, W: Semiring, F: ExpandedFst<W>, V: Visitor<'a, W, F>, A: Tr
                     }
                     state_color[tr.nextstate as usize] = DfsStateColor::Grey;
                     state_stack_next = Some(DfsState::new(fst, tr.nextstate));
-                    dfs = visitor.init_state(tr.nextstate, root);
+                    dfs = visitor.init_state(tr.nextstate, root)?;
                 }
                 DfsStateColor::Grey => {
                     dfs = visitor.back_tr(s, tr);
@@ -185,4 +187,5 @@ pub fn dfs_visit<'a, W: Semiring, F: ExpandedFst<W>, V: Visitor<'a, W, F>, A: Tr
         }
     }
     visitor.finish_visit();
+    Ok(())
 }
diff --git a/rustfst/src/algorithms/lazy/mod.rs b/rustfst/src/algorithms/lazy/mod.rs
index cd56766ed..fcee56240 100644
--- a/rustfst/src/algorithms/lazy/mod.rs
+++ b/rustfst/src/algorithms/lazy/mod.rs
@@ -1,7 +1,7 @@
 pub use self::cache::*;
 pub use fst_op::{AccessibleOpState, FstOp, SerializableOpState};
 pub use fst_op_2::FstOp2;
-pub use lazy_fst::{LazyFst, SerializableLazyFst, iterate_lazy};
+pub use lazy_fst::{iterate_lazy, LazyFst, SerializableLazyFst};
 pub use lazy_fst_2::LazyFst2;
 pub use state_table::StateTable;
 
diff --git a/rustfst/src/algorithms/queues/auto_queue.rs b/rustfst/src/algorithms/queues/auto_queue.rs
index 0330d6a45..5bce3693c 100644
--- a/rustfst/src/algorithms/queues/auto_queue.rs
+++ b/rustfst/src/algorithms/queues/auto_queue.rs
@@ -32,14 +32,14 @@ impl AutoQueue {
         if props.contains(FstProperties::TOP_SORTED) || fst.start().is_none() {
             queue = Box::new(StateOrderQueue::default());
         } else if props.contains(FstProperties::ACYCLIC) {
-            queue = Box::new(TopOrderQueue::new(fst, tr_filter));
+            queue = Box::new(TopOrderQueue::new(fst, tr_filter)?);
         } else if props.contains(FstProperties::UNWEIGHTED)
             && W::properties().contains(SemiringProperties::IDEMPOTENT)
         {
             queue = Box::new(LifoQueue::default());
         } else {
             let mut scc_visitor = SccVisitor::new(fst, true, false);
-            dfs_visit(fst, &mut scc_visitor, tr_filter, false);
+            dfs_visit(fst, &mut scc_visitor, tr_filter, false)?;
             let sccs: Vec<_> = scc_visitor
                 .scc
                 .unwrap()
diff --git a/rustfst/src/algorithms/queues/top_order_queue.rs b/rustfst/src/algorithms/queues/top_order_queue.rs
index eff2c9168..6f830ad14 100644
--- a/rustfst/src/algorithms/queues/top_order_queue.rs
+++ b/rustfst/src/algorithms/queues/top_order_queue.rs
@@ -5,6 +5,7 @@ use crate::algorithms::{Queue, QueueType};
 use crate::fst_traits::ExpandedFst;
 use crate::semirings::Semiring;
 use crate::StateId;
+use anyhow::{bail, Result};
 
 /// Topological-order queue discipline, templated on the StateId. States are
 /// ordered in the queue topologically. The FST must be acyclic.
@@ -17,13 +18,16 @@ pub struct TopOrderQueue {
 }
 
 impl TopOrderQueue {
-    pub fn new<W: Semiring, F: ExpandedFst<W>, A: TrFilter<W>>(fst: &F, tr_filter: &A) -> Self {
+    pub fn new<W: Semiring, F: ExpandedFst<W>, A: TrFilter<W>>(
+        fst: &F,
+        tr_filter: &A,
+    ) -> Result<Self> {
         let mut visitor = TopOrderVisitor::new();
-        dfs_visit(fst, &mut visitor, tr_filter, false);
+        dfs_visit(fst, &mut visitor, tr_filter, false)?;
         if !visitor.acyclic {
-            panic!("Unexpectted Acyclic FST for TopOprerQueue");
+            bail!("Unexpectted Acyclic FST for TopOprerQueue");
         }
-        Self::from_precomputed_order(visitor.order)
+        Ok(Self::from_precomputed_order(visitor.order))
     }
 
     pub fn from_precomputed_order(order: Vec<StateId>) -> Self {
diff --git a/rustfst/src/algorithms/randgen/mod.rs b/rustfst/src/algorithms/randgen/mod.rs
index 0f3cf7264..7d55afd20 100644
--- a/rustfst/src/algorithms/randgen/mod.rs
+++ b/rustfst/src/algorithms/randgen/mod.rs
@@ -46,7 +46,7 @@ pub fn randgen_with_config<
         // TODO: No need to do that if dfs_visit supports NOT expanded FST.
         let randgen_fst_static: VectorFst<_> = randgen_fst.compute()?;
         let mut visitor = RandGenVisitor::new();
-        dfs_visit(&randgen_fst_static, &mut visitor, &AnyTrFilter {}, false);
+        dfs_visit(&randgen_fst_static, &mut visitor, &AnyTrFilter {}, false)?;
         Ok(visitor.into_output_fst())
     }
 }
diff --git a/rustfst/src/algorithms/randgen/randgen_visitor.rs b/rustfst/src/algorithms/randgen/randgen_visitor.rs
index 1712961f3..b686a0e4d 100644
--- a/rustfst/src/algorithms/randgen/randgen_visitor.rs
+++ b/rustfst/src/algorithms/randgen/randgen_visitor.rs
@@ -52,8 +52,8 @@ impl<'a, W: Semiring, FI: Fst<W>, FO: MutableFst<W>> Visitor<'a, W, FI>
         self.path.clear();
     }
 
-    fn init_state(&mut self, _s: StateId, _root: StateId) -> bool {
-        true
+    fn init_state(&mut self, _s: StateId, _root: StateId) -> Result<bool> {
+        Ok(true)
     }
 
     fn tree_tr(&mut self, _s: StateId, tr: &Tr<W>) -> bool {
diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
index 0ff8e2fa3..129f793f1 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
@@ -104,7 +104,7 @@ pub(crate) fn rm_epsilon_with_internal_config<W: Semiring, F: MutableFst<W>, Q:
         states = fst.states_iter().collect();
     } else if fst_props.contains(FstProperties::ACYCLIC) {
         let mut visitor = TopOrderVisitor::new();
-        dfs_visit(fst, &mut visitor, &EpsilonTrFilter {}, false);
+        dfs_visit(fst, &mut visitor, &EpsilonTrFilter {}, false)?;
 
         states.resize(visitor.order.len(), 0);
         for i in 0..visitor.order.len() {
@@ -112,7 +112,7 @@ pub(crate) fn rm_epsilon_with_internal_config<W: Semiring, F: MutableFst<W>, Q:
         }
     } else {
         let mut visitor = SccVisitor::new(fst, true, false);
-        dfs_visit(fst, &mut visitor, &EpsilonTrFilter {}, false);
+        dfs_visit(fst, &mut visitor, &EpsilonTrFilter {}, false)?;
 
         let scc = visitor.scc.as_ref().unwrap();
 
diff --git a/rustfst/src/algorithms/rm_final_epsilon.rs b/rustfst/src/algorithms/rm_final_epsilon.rs
index e8f1f83e2..e3aee85b6 100644
--- a/rustfst/src/algorithms/rm_final_epsilon.rs
+++ b/rustfst/src/algorithms/rm_final_epsilon.rs
@@ -18,7 +18,7 @@ where
     F: MutableFst<W>,
 {
     let mut visitors = SccVisitor::new(ifst, false, true);
-    dfs_visit(ifst, &mut visitors, &AnyTrFilter {}, false);
+    dfs_visit(ifst, &mut visitors, &AnyTrFilter {}, false)?;
 
     let mut finals = HashSet::new();
 
diff --git a/rustfst/src/algorithms/top_sort.rs b/rustfst/src/algorithms/top_sort.rs
index a5cf21083..711d20113 100644
--- a/rustfst/src/algorithms/top_sort.rs
+++ b/rustfst/src/algorithms/top_sort.rs
@@ -28,8 +28,8 @@ impl TopOrderVisitor {
 impl<'a, W: Semiring, F: 'a + Fst<W>> Visitor<'a, W, F> for TopOrderVisitor {
     fn init_visit(&mut self, _fst: &'a F) {}
 
-    fn init_state(&mut self, _s: StateId, _root: StateId) -> bool {
-        true
+    fn init_state(&mut self, _s: StateId, _root: StateId) -> Result<bool> {
+        Ok(true)
     }
 
     fn tree_tr(&mut self, _s: StateId, _tr: &Tr<W>) -> bool {
@@ -79,7 +79,7 @@ where
     F: MutableFst<W>,
 {
     let mut visitor = TopOrderVisitor::new();
-    dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false);
+    dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false)?;
     if visitor.acyclic {
         state_sort(fst, &visitor.order)?;
         let props =
diff --git a/rustfst/src/algorithms/visitors/scc_visitors.rs b/rustfst/src/algorithms/visitors/scc_visitors.rs
index 07887a73b..70a010db9 100644
--- a/rustfst/src/algorithms/visitors/scc_visitors.rs
+++ b/rustfst/src/algorithms/visitors/scc_visitors.rs
@@ -1,5 +1,6 @@
 use std::marker::PhantomData;
 
+use anyhow::Result;
 use unsafe_unwrap::UnsafeUnwrap;
 
 use crate::algorithms::dfs_visit::Visitor;
@@ -62,7 +63,7 @@ impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> SccVisitor<'a, W, F> {
 impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> Visitor<'a, W, F> for SccVisitor<'a, W, F> {
     fn init_visit(&mut self, _fst: &'a F) {}
 
-    fn init_state(&mut self, s: StateId, root: StateId) -> bool {
+    fn init_state(&mut self, s: StateId, root: StateId) -> Result<bool> {
         self.scc_stack.push(s);
         let s = s as usize;
         self.dfnumber[s] = self.nstates as i32;
@@ -80,7 +81,7 @@ impl<'a, W: Semiring, F: 'a + ExpandedFst<W>> Visitor<'a, W, F> for SccVisitor<'
             self.props &= !FstProperties::ACCESSIBLE;
         }
         self.nstates += 1;
-        true
+        Ok(true)
     }
 
     fn tree_tr(&mut self, _s: StateId, _tr: &Tr<W>) -> bool {
diff --git a/rustfst/src/fst_properties/compute_fst_properties.rs b/rustfst/src/fst_properties/compute_fst_properties.rs
index d8558c430..b6400915f 100644
--- a/rustfst/src/fst_properties/compute_fst_properties.rs
+++ b/rustfst/src/fst_properties/compute_fst_properties.rs
@@ -43,7 +43,7 @@ pub fn compute_fst_properties<W: Semiring, F: ExpandedFst<W>>(
         .intersects(dfs_props | FstProperties::WEIGHTED_CYCLES | FstProperties::UNWEIGHTED_CYCLES)
     {
         let mut visitor = SccVisitor::new(fst, true, true);
-        dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false);
+        dfs_visit(fst, &mut visitor, &AnyTrFilter {}, false)?;
 
         // Retrieves props computed in the DFS.
         comp_props |= dfs_props & visitor.props;
diff --git a/rustfst/src/fst_traits/expanded_fst.rs b/rustfst/src/fst_traits/expanded_fst.rs
index 17720bfc3..d2f069c84 100644
--- a/rustfst/src/fst_traits/expanded_fst.rs
+++ b/rustfst/src/fst_traits/expanded_fst.rs
@@ -60,7 +60,16 @@ pub trait ExpandedFst<W: Semiring>: Fst<W> + Clone + PartialEq + FstIntoIterator
                     || tr1.olabel != tr2.olabel
                     || tr1.nextstate != tr2.nextstate
                 {
-                    println!("State {} Ilabel {} {} Olabel {} {} Nextstate {} {}", state, tr1.ilabel, tr2.ilabel, tr1.olabel, tr2.olabel, tr1.nextstate, tr2.nextstate);
+                    println!(
+                        "State {} Ilabel {} {} Olabel {} {} Nextstate {} {}",
+                        state,
+                        tr1.ilabel,
+                        tr2.ilabel,
+                        tr1.olabel,
+                        tr2.olabel,
+                        tr1.nextstate,
+                        tr2.nextstate
+                    );
                     return false;
                 }
 
diff --git a/rustfst/src/tests_openfst/algorithms/compose.rs b/rustfst/src/tests_openfst/algorithms/compose.rs
index 34e8f1310..3a3ceb29f 100644
--- a/rustfst/src/tests_openfst/algorithms/compose.rs
+++ b/rustfst/src/tests_openfst/algorithms/compose.rs
@@ -24,12 +24,12 @@ use crate::algorithms::lazy::SimpleHashMapCache;
 use crate::algorithms::{tr_compares::ILabelCompare, tr_sort};
 use crate::fst_impls::VectorFst;
 use crate::fst_traits::{CoreFst, ExpandedFst, Fst, MutableFst, SerializableFst};
+use crate::prelude::lazy::SerializableLazyFst;
 use crate::semirings::{SerializableSemiring, WeaklyDivisibleSemiring, WeightQuantize};
 use crate::tests_openfst::utils::test_eq_fst;
 use crate::tests_openfst::FstTestData;
 use crate::{StateId, SymbolTable, Trs};
 use std::cmp::max;
-use crate::prelude::lazy::SerializableLazyFst;
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct ComposeOperationResult {

From 1f250f879c74721997026a164b7a44cb65e7a90b Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 3 Jan 2023 18:02:41 +0100
Subject: [PATCH 5/9] Add logs

---
 rustfst/src/algorithms/compose/compose_fst.rs |  4 +--
 .../lookahead_relabel_fst.rs                  |  2 +-
 rustfst/src/algorithms/lazy/lazy_fst.rs       | 30 ++++++++++++++++---
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/rustfst/src/algorithms/compose/compose_fst.rs b/rustfst/src/algorithms/compose/compose_fst.rs
index 72c411ba6..d9a48c832 100644
--- a/rustfst/src/algorithms/compose/compose_fst.rs
+++ b/rustfst/src/algorithms/compose/compose_fst.rs
@@ -163,8 +163,8 @@ where
     /// Turns the Lazy FST into a static one.
     pub fn compute<F: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F> {
         // Small trick to make sure that both FSTs are fully expanded.
-        iterate_lazy(self.0.op.fst1.borrow())?;
-        iterate_lazy(self.0.op.fst2.borrow())?;
+        // iterate_lazy(self.0.op.fst1.borrow())?;
+        // iterate_lazy(self.0.op.fst2.borrow())?;
         self.0.compute()
     }
 }
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
index a574f2cfa..ad296471d 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -37,7 +37,7 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFst<W, F, B> {
     pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
         let underneath_fst = self.0.op.fst.borrow();
         // Trick to compute the underneath lazy fst
-        iterate_lazy(underneath_fst)?;
+        // iterate_lazy(underneath_fst)?;
         self.0.compute()
     }
 }
diff --git a/rustfst/src/algorithms/lazy/lazy_fst.rs b/rustfst/src/algorithms/lazy/lazy_fst.rs
index ed1ea3835..0934c19d3 100644
--- a/rustfst/src/algorithms/lazy/lazy_fst.rs
+++ b/rustfst/src/algorithms/lazy/lazy_fst.rs
@@ -8,7 +8,6 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use itertools::izip;
-use unsafe_unwrap::UnsafeUnwrap;
 
 use crate::algorithms::lazy::cache::CacheStatus;
 use crate::algorithms::lazy::fst_op::{AccessibleOpState, FstOp, SerializableOpState};
@@ -57,7 +56,15 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     unsafe fn final_weight_unchecked(&self, state_id: StateId) -> Option<W> {
-        self.final_weight(state_id).unsafe_unwrap()
+        match self.final_weight(state_id) {
+            Result::Err(e) => panic!(
+                "`final_weight_unchecked` failed for state {} which should never happen. \
+                Most likely an issue in the implementation of the corresponding LazyFst. \
+                Underneath Error is {:?}\
+                ", state_id, e
+            ),
+            Ok(v) => v
+        }
     }
 
     fn num_trs(&self, s: StateId) -> Result<usize> {
@@ -69,7 +76,14 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
-        self.cache.num_trs(s).unsafe_unwrap()
+        match self.cache.num_trs(s) {
+            None => panic!(
+                "`num_trs_unchecked` failed which should never happen. \
+                `self.cache.num_trs` returned None indicating that the state {} doesn't exist.
+                ", s
+            ),
+            Some(v) => v
+        }
     }
 
     fn get_trs(&self, state_id: StateId) -> Result<Self::TRS> {
@@ -84,7 +98,15 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     unsafe fn get_trs_unchecked(&self, state_id: StateId) -> Self::TRS {
-        self.get_trs(state_id).unsafe_unwrap()
+        match self.get_trs(state_id) {
+            Result::Err(e) => panic!(
+                "`get_trs_unchecked` failed for state {} which should never happen. \
+                Most likely an issue in the implementation of the corresponding LazyFst. \
+                Underneath Error is {:?}\
+                ", state_id, e
+            ),
+            Result::Ok(v) => v
+        }
     }
 
     fn properties(&self) -> FstProperties {

From 4f5d68f447ce667448b729eb6a7db5115d14065a Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 3 Jan 2023 18:26:13 +0100
Subject: [PATCH 6/9] fmt

---
 .../lookahead_relabel/lookahead_relabel_fst.rs    |  2 +-
 rustfst/src/algorithms/lazy/lazy_fst.rs           | 15 +++++++++------
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
index ad296471d..3c16b1e07 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -1,5 +1,5 @@
 use crate::algorithms::compose::lookahead_relabel::lookahead_relabel_fst_op::LookaheadRelabelFstOp;
-use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache, iterate_lazy};
+use crate::algorithms::lazy::{iterate_lazy, LazyFst, SimpleHashMapCache};
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
 use crate::prelude::compose::LabelReachableData;
diff --git a/rustfst/src/algorithms/lazy/lazy_fst.rs b/rustfst/src/algorithms/lazy/lazy_fst.rs
index 0934c19d3..06a258e9f 100644
--- a/rustfst/src/algorithms/lazy/lazy_fst.rs
+++ b/rustfst/src/algorithms/lazy/lazy_fst.rs
@@ -61,9 +61,10 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
                 "`final_weight_unchecked` failed for state {} which should never happen. \
                 Most likely an issue in the implementation of the corresponding LazyFst. \
                 Underneath Error is {:?}\
-                ", state_id, e
+                ",
+                state_id, e
             ),
-            Ok(v) => v
+            Ok(v) => v,
         }
     }
 
@@ -80,9 +81,10 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
             None => panic!(
                 "`num_trs_unchecked` failed which should never happen. \
                 `self.cache.num_trs` returned None indicating that the state {} doesn't exist.
-                ", s
+                ",
+                s
             ),
-            Some(v) => v
+            Some(v) => v,
         }
     }
 
@@ -103,9 +105,10 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
                 "`get_trs_unchecked` failed for state {} which should never happen. \
                 Most likely an issue in the implementation of the corresponding LazyFst. \
                 Underneath Error is {:?}\
-                ", state_id, e
+                ",
+                state_id, e
             ),
-            Result::Ok(v) => v
+            Result::Ok(v) => v,
         }
     }
 

From d36c2082b57f4d4a2603418b346d0faca2eb7488 Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Thu, 5 Jan 2023 15:27:15 +0100
Subject: [PATCH 7/9] Save work

---
 rustfst/src/algorithms/compose/compose_fst.rs |  1 -
 .../src/algorithms/compose/compose_fst_op.rs  |  3 +-
 .../lookahead_relabel_fst.rs                  |  4 +-
 rustfst/src/algorithms/lazy/lazy_fst.rs       | 62 +++++++++----------
 rustfst/src/algorithms/lazy/mod.rs            |  2 +-
 5 files changed, 36 insertions(+), 36 deletions(-)

diff --git a/rustfst/src/algorithms/compose/compose_fst.rs b/rustfst/src/algorithms/compose/compose_fst.rs
index d9a48c832..62b56aad7 100644
--- a/rustfst/src/algorithms/compose/compose_fst.rs
+++ b/rustfst/src/algorithms/compose/compose_fst.rs
@@ -12,7 +12,6 @@ use crate::algorithms::compose::matchers::{GenericMatcher, Matcher};
 use crate::algorithms::compose::{
     ComposeFstOp, ComposeFstOpOptions, ComposeFstOpState, ComposeStateTuple,
 };
-use crate::algorithms::lazy::iterate_lazy;
 use crate::algorithms::lazy::{
     FstCache, LazyFst, SerializableCache, SerializableLazyFst, SimpleVecCache,
 };
diff --git a/rustfst/src/algorithms/compose/compose_fst_op.rs b/rustfst/src/algorithms/compose/compose_fst_op.rs
index bbbd3f08d..ca5e80954 100644
--- a/rustfst/src/algorithms/compose/compose_fst_op.rs
+++ b/rustfst/src/algorithms/compose/compose_fst_op.rs
@@ -389,11 +389,12 @@ where
     fn compute_start(&self) -> Result<Option<StateId>> {
         let compose_filter = self.compose_filter_builder.build()?;
         let s1 = self.fst1.borrow().start();
+        // Let's put it here to force the fst2 to have its start state computed in the case of a Lazy Fst.
+        let s2 = self.fst2.borrow().start();
         if s1.is_none() {
             return Ok(None);
         }
         let s1 = s1.unwrap();
-        let s2 = self.fst2.borrow().start();
         if s2.is_none() {
             return Ok(None);
         }
diff --git a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
index 3c16b1e07..4b7d5f618 100644
--- a/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
+++ b/rustfst/src/algorithms/compose/lookahead_relabel/lookahead_relabel_fst.rs
@@ -1,5 +1,5 @@
 use crate::algorithms::compose::lookahead_relabel::lookahead_relabel_fst_op::LookaheadRelabelFstOp;
-use crate::algorithms::lazy::{iterate_lazy, LazyFst, SimpleHashMapCache};
+use crate::algorithms::lazy::{LazyFst, SimpleHashMapCache};
 use crate::fst_properties::FstProperties;
 use crate::fst_traits::{AllocableFst, CoreFst, Fst, FstIterator, MutableFst, StateIterator};
 use crate::prelude::compose::LabelReachableData;
@@ -35,7 +35,7 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> LookaheadRelabelFst<W, F, B> {
     }
 
     pub fn compute<F2: MutableFst<W> + AllocableFst<W>>(&self) -> Result<F2> {
-        let underneath_fst = self.0.op.fst.borrow();
+        // let underneath_fst = self.0.op.fst.borrow();
         // Trick to compute the underneath lazy fst
         // iterate_lazy(underneath_fst)?;
         self.0.compute()
diff --git a/rustfst/src/algorithms/lazy/lazy_fst.rs b/rustfst/src/algorithms/lazy/lazy_fst.rs
index 06a258e9f..480d994ab 100644
--- a/rustfst/src/algorithms/lazy/lazy_fst.rs
+++ b/rustfst/src/algorithms/lazy/lazy_fst.rs
@@ -69,11 +69,11 @@ impl<W: Semiring, Op: FstOp<W>, Cache: FstCache<W>> CoreFst<W> for LazyFst<W, Op
     }
 
     fn num_trs(&self, s: StateId) -> Result<usize> {
-        // let trs = self.get_trs(s)?;
-        // Ok(trs.trs().len())
-        self.cache
-            .num_trs(s)
-            .ok_or_else(|| format_err!("NT State {:?} doesn't exist", s))
+        let trs = self.get_trs(s)?;
+        Ok(trs.trs().len())
+        // self.cache
+        //     .num_trs(s)
+        //     .ok_or_else(|| format_err!("`num_trs` failed -> State {:?} doesn't exist", s))
     }
 
     unsafe fn num_trs_unchecked(&self, s: StateId) -> usize {
@@ -298,32 +298,32 @@ where
 }
 
 /// Turns the Lazy FST into a static one.
-pub fn iterate_lazy<W: Semiring, F: Fst<W>>(fst: &F) -> Result<()> {
-    let start_state = fst.start();
-    let start_state = match start_state {
-        Some(s) => s,
-        None => return Ok(()),
-    };
-    let mut queue = VecDeque::new();
-    let mut visited_states = vec![];
-    visited_states.resize(start_state as usize + 1, false);
-    visited_states[start_state as usize] = true;
-    queue.push_back(start_state);
-    while let Some(s) = queue.pop_front() {
-        let trs_owner = fst.get_trs(s)?;
-        for tr in trs_owner.trs() {
-            if (tr.nextstate as usize) >= visited_states.len() {
-                visited_states.resize(tr.nextstate as usize + 1, false);
-            }
-            if !visited_states[tr.nextstate as usize] {
-                queue.push_back(tr.nextstate);
-                visited_states[tr.nextstate as usize] = true;
-            }
-        }
-        fst.final_weight(s)?;
-    }
-    Ok(())
-}
+// pub fn iterate_lazy<W: Semiring, F: Fst<W>>(fst: &F) -> Result<()> {
+//     let start_state = fst.start();
+//     let start_state = match start_state {
+//         Some(s) => s,
+//         None => return Ok(()),
+//     };
+//     let mut queue = VecDeque::new();
+//     let mut visited_states = vec![];
+//     visited_states.resize(start_state as usize + 1, false);
+//     visited_states[start_state as usize] = true;
+//     queue.push_back(start_state);
+//     while let Some(s) = queue.pop_front() {
+//         let trs_owner = fst.get_trs(s)?;
+//         for tr in trs_owner.trs() {
+//             if (tr.nextstate as usize) >= visited_states.len() {
+//                 visited_states.resize(tr.nextstate as usize + 1, false);
+//             }
+//             if !visited_states[tr.nextstate as usize] {
+//                 queue.push_back(tr.nextstate);
+//                 visited_states[tr.nextstate as usize] = true;
+//             }
+//         }
+//         fst.final_weight(s)?;
+//     }
+//     Ok(())
+// }
 
 impl<W, Op, Cache> SerializableLazyFst for LazyFst<W, Op, Cache>
 where
diff --git a/rustfst/src/algorithms/lazy/mod.rs b/rustfst/src/algorithms/lazy/mod.rs
index fcee56240..68368fab8 100644
--- a/rustfst/src/algorithms/lazy/mod.rs
+++ b/rustfst/src/algorithms/lazy/mod.rs
@@ -1,7 +1,7 @@
 pub use self::cache::*;
 pub use fst_op::{AccessibleOpState, FstOp, SerializableOpState};
 pub use fst_op_2::FstOp2;
-pub use lazy_fst::{iterate_lazy, LazyFst, SerializableLazyFst};
+pub use lazy_fst::{LazyFst, SerializableLazyFst};
 pub use lazy_fst_2::LazyFst2;
 pub use state_table::StateTable;
 

From 62025d5607ebf594fb5dfa18f36acfa6451bcb1a Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 31 Jan 2023 16:56:38 +0100
Subject: [PATCH 8/9] RmEpsilon on lazy fsts

---
 .../algorithms/rm_epsilon/rm_epsilon_fst.rs   | 17 ++---
 .../algorithms/rm_epsilon/rm_epsilon_op.rs    | 26 +++----
 .../algorithms/rm_epsilon/rm_epsilon_state.rs | 43 ++++++-----
 .../rm_epsilon/rm_epsilon_static.rs           | 21 ++++--
 rustfst/src/algorithms/shortest_distance.rs   | 72 ++++++++++++-------
 5 files changed, 106 insertions(+), 73 deletions(-)

diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
index 0e60a39f8..a2214a68b 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
@@ -7,7 +7,8 @@ use anyhow::Result;
 use crate::algorithms::lazy::{LazyFst2, SimpleHashMapCache};
 use crate::algorithms::rm_epsilon::rm_epsilon_op::RmEpsilonOp;
 use crate::fst_properties::FstProperties;
-use crate::fst_traits::{CoreFst, Fst, FstIterator, MutableFst, StateIterator};
+use crate::fst_traits::{CoreFst, Fst, FstIterator, StateIterator};
+use crate::prelude::MutableFst;
 use crate::{Semiring, StateId, SymbolTable, TrsVec};
 
 type InnerLazyFst<W, F, B> = LazyFst2<W, RmEpsilonOp<W, F, B>, SimpleHashMapCache<W>>;
@@ -17,12 +18,12 @@ type InnerLazyFst<W, F, B> = LazyFst2<W, RmEpsilonOp<W, F, B>, SimpleHashMapCach
 /// States and transitions will be added as necessary. The algorithm is a
 /// generalization to arbitrary weights of the second step of the input
 /// epsilon-normalization algorithm. This version is a Delayed FST.
-pub struct RmEpsilonFst<W: Semiring, F: MutableFst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
+pub struct RmEpsilonFst<W: Semiring, F: Fst<W>, B: Borrow<F>>(InnerLazyFst<W, F, B>);
 
 impl<W, F, B> CoreFst<W> for RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W>,
+    F: Fst<W>,
     B: Borrow<F>,
 {
     type TRS = TrsVec<W>;
@@ -71,7 +72,7 @@ where
 impl<'a, W, F, B> StateIterator<'a> for RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W> + 'a,
+    F: Fst<W> + 'a,
     B: Borrow<F> + 'a,
 {
     type Iter = <InnerLazyFst<W, F, B> as StateIterator<'a>>::Iter;
@@ -84,7 +85,7 @@ where
 impl<'a, W, F, B> FstIterator<'a, W> for RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W> + 'a,
+    F: Fst<W> + 'a,
     B: Borrow<F> + 'a,
 {
     type FstIter = <InnerLazyFst<W, F, B> as FstIterator<'a, W>>::FstIter;
@@ -97,7 +98,7 @@ where
 impl<W, F, B> Fst<W> for RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W> + 'static,
+    F: Fst<W> + 'static,
     B: Borrow<F> + 'static,
 {
     fn input_symbols(&self) -> Option<&Arc<SymbolTable>> {
@@ -128,7 +129,7 @@ where
 impl<W, F, B> Debug for RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W>,
+    F: Fst<W>,
     B: Borrow<F>,
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -139,7 +140,7 @@ where
 impl<'a, W, F, B> RmEpsilonFst<W, F, B>
 where
     W: Semiring,
-    F: MutableFst<W>,
+    F: Fst<W>,
     B: Borrow<F>,
 {
     pub fn new(fst: B) -> Result<Self> {
diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
index 68359080d..a31db70c1 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
@@ -7,23 +7,20 @@ use crate::algorithms::queues::FifoQueue;
 use crate::algorithms::rm_epsilon::{RmEpsilonInternalConfig, RmEpsilonState};
 use crate::fst_properties::mutable_properties::rmepsilon_properties;
 use crate::fst_properties::FstProperties;
-use crate::fst_traits::MutableFst;
+use crate::fst_traits::Fst;
 use crate::semirings::Semiring;
 use crate::{StateId, TrsVec};
 use itertools::Itertools;
 use std::cell::RefCell;
-use std::marker::PhantomData;
 use std::sync::Arc;
 
 #[derive(Clone)]
-pub struct RmEpsilonOp<W: Semiring, F: MutableFst<W>, B: Borrow<F>> {
-    rmeps_state: RefCell<RmEpsilonState<W, FifoQueue>>,
+pub struct RmEpsilonOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
+    rmeps_state: RefCell<RmEpsilonState<W, FifoQueue, F, B>>,
     properties: FstProperties,
-    ghost: PhantomData<F>,
-    fst: B,
 }
 
-impl<W: Semiring, F: MutableFst<W>, B: Borrow<F>> std::fmt::Debug for RmEpsilonOp<W, F, B> {
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> std::fmt::Debug for RmEpsilonOp<W, F, B> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
@@ -39,31 +36,26 @@ impl<W: Semiring, F: MutableFst<W>, B: Borrow<F>> std::fmt::Debug for RmEpsilonO
 //     }
 // }
 
-impl<W: Semiring, F: MutableFst<W>, B: Borrow<F>> RmEpsilonOp<W, F, B> {
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RmEpsilonOp<W, F, B> {
     pub fn new(fst: B) -> Self {
         let properties = rmepsilon_properties(fst.borrow().properties(), true);
         Self {
             properties,
             rmeps_state: RefCell::new(RmEpsilonState::new(
-                fst.borrow().num_states(),
+                fst,
                 RmEpsilonInternalConfig::new_with_default(FifoQueue::default()),
             )),
-            fst,
-            ghost: PhantomData,
         }
     }
 }
 
-impl<W: Semiring, F: MutableFst<W>, B: Borrow<F>> FstOp2<W> for RmEpsilonOp<W, F, B> {
+impl<W: Semiring, F: Fst<W>, B: Borrow<F>> FstOp2<W> for RmEpsilonOp<W, F, B> {
     fn compute_start(&self) -> Result<Option<StateId>> {
-        Ok(self.fst.borrow().start())
+        Ok(self.rmeps_state.borrow().fst().borrow().start())
     }
 
     fn compute_trs_and_final_weight(&self, state: StateId) -> Result<(TrsVec<W>, Option<W>)> {
-        let (trs, final_weight) = self
-            .rmeps_state
-            .borrow_mut()
-            .expand::<F, _>(state, self.fst.borrow())?;
+        let (trs, final_weight) = self.rmeps_state.borrow_mut().expand(state)?;
         let zero = W::zero();
 
         let trs = trs.into_iter().rev().collect_vec();
diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_state.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_state.rs
index fa59eac74..e10645f3b 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_state.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_state.rs
@@ -8,45 +8,48 @@ use crate::algorithms::rm_epsilon::{Element, RmEpsilonInternalConfig};
 use crate::algorithms::shortest_distance::ShortestDistanceState;
 use crate::algorithms::tr_filters::{EpsilonTrFilter, TrFilter};
 use crate::algorithms::Queue;
-use crate::fst_traits::ExpandedFst;
+use crate::prelude::Fst;
 use crate::semirings::Semiring;
 use crate::{StateId, Tr, Trs};
+use std::marker::PhantomData;
 
 #[derive(Clone)]
-pub(crate) struct RmEpsilonState<W: Semiring, Q: Queue> {
+pub(crate) struct RmEpsilonState<W: Semiring, Q: Queue, F: Fst<W>, B: Borrow<F>> {
     pub visited: Vec<bool>,
     pub visited_states: Vec<StateId>,
     pub element_map: HashMap<Element, (StateId, usize)>,
     pub expand_id: StateId,
-    pub sd_state: ShortestDistanceState<W, Q, EpsilonTrFilter>,
+    pub sd_state: ShortestDistanceState<W, Q, EpsilonTrFilter, F, B>,
+    ghost: PhantomData<F>,
 }
 
-impl<W: Semiring, Q: Queue> std::fmt::Debug for RmEpsilonState<W, Q> {
+impl<W: Semiring, Q: Queue, F: Fst<W>, B: Borrow<F>> std::fmt::Debug
+    for RmEpsilonState<W, Q, F, B>
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "RmEpsilonState {{ visited : {:?}, visited_states : {:?}, element_map : {:?}, expand_id : {:?}, sd_state : {:?} }}",
                self.visited, self.visited_states, self.element_map, self.expand_id, self.sd_state)
     }
 }
 
-impl<W: Semiring, Q: Queue> RmEpsilonState<W, Q> {
-    pub fn new(fst_num_states: usize, opts: RmEpsilonInternalConfig<W, Q>) -> Self {
+impl<W: Semiring, Q: Queue, F: Fst<W>, B: Borrow<F>> RmEpsilonState<W, Q, F, B> {
+    pub fn new(fst: B, opts: RmEpsilonInternalConfig<W, Q>) -> Self {
         Self {
-            sd_state: ShortestDistanceState::new_from_config(fst_num_states, opts.sd_opts, true),
+            sd_state: ShortestDistanceState::new_from_config(opts.sd_opts, true, fst),
             visited: vec![],
             visited_states: vec![],
             element_map: HashMap::new(),
             expand_id: 0,
+            ghost: PhantomData,
         }
     }
 
-    pub fn expand<F: ExpandedFst<W>, B: Borrow<F>>(
-        &mut self,
-        source: StateId,
-        fst: B,
-    ) -> Result<(Vec<Tr<W>>, W)> {
-        let distance = self
-            .sd_state
-            .shortest_distance::<F, _>(Some(source), fst.borrow())?;
+    pub fn fst(&self) -> &F {
+        self.sd_state.fst()
+    }
+
+    pub fn expand(&mut self, source: StateId) -> Result<(Vec<Tr<W>>, W)> {
+        let distance = self.sd_state.shortest_distance(Some(source))?;
 
         let tr_filter = EpsilonTrFilter {};
 
@@ -63,7 +66,7 @@ impl<W: Semiring, Q: Queue> RmEpsilonState<W, Q> {
             }
             self.visited[state as usize] = true;
             self.visited_states.push(state);
-            for tr in fst.borrow().get_trs(state)?.trs() {
+            for tr in self.sd_state.fst().get_trs(state)?.trs() {
                 // TODO: Remove this clone
                 let mut tr = tr.clone();
                 tr.weight = distance[state as usize].times(&tr.weight)?;
@@ -104,8 +107,12 @@ impl<W: Semiring, Q: Queue> RmEpsilonState<W, Q> {
                 }
             }
             final_weight.plus_assign(
-                distance[state as usize]
-                    .times(fst.borrow().final_weight(state)?.unwrap_or_else(W::zero))?,
+                distance[state as usize].times(
+                    self.sd_state
+                        .fst()
+                        .final_weight(state)?
+                        .unwrap_or_else(W::zero),
+                )?,
             )?;
         }
 
diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
index 129f793f1..a3bc8a21e 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_static.rs
@@ -134,7 +134,7 @@ pub(crate) fn rm_epsilon_with_internal_config<W: Semiring, F: MutableFst<W>, Q:
         }
     }
 
-    let mut rmeps_state = RmEpsilonState::new(fst.num_states(), opts);
+    let mut rmeps_state: RmEpsilonState<_, _, F, _> = RmEpsilonState::new(&mut *fst, opts);
     let zero = W::zero();
 
     for state in states.into_iter().rev() {
@@ -143,15 +143,24 @@ pub(crate) fn rm_epsilon_with_internal_config<W: Semiring, F: MutableFst<W>, Q:
         {
             continue;
         }
-        let (trs, final_weight) = rmeps_state.expand::<F, _>(state, &*fst)?;
+        let (trs, final_weight) = rmeps_state.expand(state)?;
 
         unsafe {
-            fst.pop_trs_unchecked(state);
-            fst.set_trs_unchecked(state, trs.into_iter().rev().collect());
+            rmeps_state.sd_state.fst.pop_trs_unchecked(state);
+            rmeps_state
+                .sd_state
+                .fst
+                .set_trs_unchecked(state, trs.into_iter().rev().collect());
             if final_weight != zero {
-                fst.set_final_unchecked(state, final_weight);
+                rmeps_state
+                    .sd_state
+                    .fst
+                    .set_final_unchecked(state, final_weight);
             } else {
-                fst.delete_final_weight_unchecked(state);
+                rmeps_state
+                    .sd_state
+                    .fst
+                    .delete_final_weight_unchecked(state);
             }
         }
     }
diff --git a/rustfst/src/algorithms/shortest_distance.rs b/rustfst/src/algorithms/shortest_distance.rs
index 5400ef28e..c2811dc40 100644
--- a/rustfst/src/algorithms/shortest_distance.rs
+++ b/rustfst/src/algorithms/shortest_distance.rs
@@ -6,7 +6,7 @@ use crate::algorithms::queues::AutoQueue;
 use crate::algorithms::tr_filters::{AnyTrFilter, TrFilter};
 use crate::algorithms::Queue;
 use crate::fst_impls::VectorFst;
-use crate::fst_traits::{ExpandedFst, MutableFst};
+use crate::fst_traits::{ExpandedFst, Fst, MutableFst};
 use crate::semirings::{ReverseBack, Semiring, SemiringProperties, WeightQuantize};
 use crate::{StateId, Trs, KSHORTESTDELTA};
 use std::borrow::Borrow;
@@ -45,7 +45,14 @@ impl<W: Semiring, Q: Queue, A: TrFilter<W>> ShortestDistanceInternalConfig<W, Q,
 }
 
 #[derive(Clone)]
-pub(crate) struct ShortestDistanceState<W: Semiring, Q: Queue, A: TrFilter<W>> {
+pub(crate) struct ShortestDistanceState<W, Q, A, F, B>
+where
+    W: Semiring,
+    Q: Queue,
+    A: TrFilter<W>,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
     state_queue: Q,
     tr_filter: A,
     first_path: bool,
@@ -57,9 +64,18 @@ pub(crate) struct ShortestDistanceState<W: Semiring, Q: Queue, A: TrFilter<W>> {
     retain: bool,
     source_id: usize,
     delta: f32,
+    pub(crate) fst: B,
+    ghost: PhantomData<F>,
 }
 
-impl<W: Semiring, Q: Queue, A: TrFilter<W>> std::fmt::Debug for ShortestDistanceState<W, Q, A> {
+impl<W, Q, A, F, B> std::fmt::Debug for ShortestDistanceState<W, Q, A, F, B>
+where
+    W: Semiring,
+    Q: Queue,
+    A: TrFilter<W>,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "ShortestDistanceState {{ ")?;
         write!(f, "state_queue : {:?}, ", self.state_queue)?;
@@ -97,44 +113,57 @@ macro_rules! ensure_source_index_is_valid {
     };
 }
 
-impl<W: Semiring, Q: Queue, A: TrFilter<W>> ShortestDistanceState<W, Q, A> {
+impl<W, Q, A, F, B> ShortestDistanceState<W, Q, A, F, B>
+where
+    W: Semiring,
+    Q: Queue,
+    A: TrFilter<W>,
+    F: Fst<W>,
+    B: Borrow<F>,
+{
     pub fn new(
-        fst_num_states: usize,
         state_queue: Q,
         tr_filter: A,
         first_path: bool,
         retain: bool,
         delta: f32,
+        fst: B,
     ) -> Self {
         Self {
             state_queue,
             tr_filter,
             first_path,
-            distance: Vec::with_capacity(fst_num_states),
-            enqueued: Vec::with_capacity(fst_num_states),
-            adder: Vec::with_capacity(fst_num_states),
-            radder: Vec::with_capacity(fst_num_states),
-            sources: Vec::with_capacity(fst_num_states),
+            distance: vec![],
+            enqueued: vec![],
+            adder: vec![],
+            radder: vec![],
+            sources: vec![],
             source_id: 0,
             retain,
             delta,
+            fst,
+            ghost: PhantomData,
         }
     }
     pub fn new_from_config(
-        fst_num_states: usize,
         opts: ShortestDistanceInternalConfig<W, Q, A>,
         retain: bool,
+        fst: B,
     ) -> Self {
         Self::new(
-            fst_num_states,
             opts.state_queue,
             opts.tr_filter,
             opts.first_path,
             retain,
             opts.delta,
+            fst,
         )
     }
 
+    pub fn fst(&self) -> &F {
+        self.fst.borrow()
+    }
+
     fn ensure_distance_index_is_valid(&mut self, index: usize) {
         while self.distance.len() <= index {
             self.distance.push(W::zero());
@@ -150,12 +179,8 @@ impl<W: Semiring, Q: Queue, A: TrFilter<W>> ShortestDistanceState<W, Q, A> {
         }
     }
 
-    pub fn shortest_distance<F: ExpandedFst<W>, B: Borrow<F>>(
-        &mut self,
-        source: Option<StateId>,
-        fst: B,
-    ) -> Result<Vec<W>> {
-        let start_state = match fst.borrow().start() {
+    pub fn shortest_distance(&mut self, source: Option<StateId>) -> Result<Vec<W>> {
+        let start_state = match self.fst.borrow().start() {
             Some(start_state) => start_state,
             None => return Ok(vec![]),
         };
@@ -189,13 +214,13 @@ impl<W: Semiring, Q: Queue, A: TrFilter<W>> ShortestDistanceState<W, Q, A> {
             let state = self.state_queue.head().unwrap() as usize;
             self.state_queue.dequeue();
             //            self.ensure_distance_index_is_valid(state);
-            if self.first_path && fst.borrow().is_final(state as StateId)? {
+            if self.first_path && self.fst.borrow().is_final(state as StateId)? {
                 break;
             }
             self.enqueued[state] = false;
             let r = self.radder[state].clone();
             self.radder[state] = W::zero();
-            for tr in fst.borrow().get_trs(state as StateId)?.trs() {
+            for tr in self.fst.borrow().get_trs(state as StateId)?.trs() {
                 let nextstate = tr.nextstate as usize;
                 if !self.tr_filter.keep(tr) {
                     continue;
@@ -241,15 +266,14 @@ pub(crate) fn shortest_distance_with_internal_config<
     W: Semiring,
     Q: Queue,
     A: TrFilter<W>,
-    F: ExpandedFst<W>,
+    F: Fst<W>,
 >(
     fst: &F,
     opts: ShortestDistanceInternalConfig<W, Q, A>,
 ) -> Result<Vec<W>> {
     let source = opts.source;
-    let mut sd_state =
-        ShortestDistanceState::<_, _, _>::new_from_config(fst.num_states(), opts, false);
-    sd_state.shortest_distance::<F, _>(source, fst)
+    let mut sd_state = ShortestDistanceState::<_, _, _, F, _>::new_from_config(opts, false, fst);
+    sd_state.shortest_distance(source)
 }
 
 #[derive(Debug, Clone, Copy, PartialOrd, PartialEq)]

From 2cac950c35aec36cab3725d1d3f7f7ca71064070 Mon Sep 17 00:00:00 2001
From: Garvys <alexandre.caulier@protonmail.com>
Date: Tue, 31 Jan 2023 18:12:31 +0100
Subject: [PATCH 9/9] Make RmEpsilonFst Sync

---
 .../src/algorithms/rm_epsilon/rm_epsilon_fst.rs | 14 ++++++++++++++
 .../src/algorithms/rm_epsilon/rm_epsilon_op.rs  | 17 ++++++++++-------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
index a2214a68b..ab3d5a8cf 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_fst.rs
@@ -157,3 +157,17 @@ where
         self.0.compute()
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::fst_impls::VectorFst;
+    use crate::semirings::TropicalWeight;
+
+    use super::*;
+
+    #[test]
+    fn test_rmepsilon_fst_sync() {
+        fn is_sync<T: Sync>() {}
+        is_sync::<RmEpsilonFst<TropicalWeight, VectorFst<_>, VectorFst<_>>>();
+    }
+}
diff --git a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
index a31db70c1..746a11d89 100644
--- a/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
+++ b/rustfst/src/algorithms/rm_epsilon/rm_epsilon_op.rs
@@ -11,12 +11,11 @@ use crate::fst_traits::Fst;
 use crate::semirings::Semiring;
 use crate::{StateId, TrsVec};
 use itertools::Itertools;
-use std::cell::RefCell;
-use std::sync::Arc;
+use std::ops::{Deref, DerefMut};
+use std::sync::{Arc, Mutex};
 
-#[derive(Clone)]
 pub struct RmEpsilonOp<W: Semiring, F: Fst<W>, B: Borrow<F>> {
-    rmeps_state: RefCell<RmEpsilonState<W, FifoQueue, F, B>>,
+    rmeps_state: Mutex<RmEpsilonState<W, FifoQueue, F, B>>,
     properties: FstProperties,
 }
 
@@ -41,7 +40,7 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RmEpsilonOp<W, F, B> {
         let properties = rmepsilon_properties(fst.borrow().properties(), true);
         Self {
             properties,
-            rmeps_state: RefCell::new(RmEpsilonState::new(
+            rmeps_state: Mutex::new(RmEpsilonState::new(
                 fst,
                 RmEpsilonInternalConfig::new_with_default(FifoQueue::default()),
             )),
@@ -51,11 +50,15 @@ impl<W: Semiring, F: Fst<W>, B: Borrow<F>> RmEpsilonOp<W, F, B> {
 
 impl<W: Semiring, F: Fst<W>, B: Borrow<F>> FstOp2<W> for RmEpsilonOp<W, F, B> {
     fn compute_start(&self) -> Result<Option<StateId>> {
-        Ok(self.rmeps_state.borrow().fst().borrow().start())
+        let mutex = self.rmeps_state.lock().unwrap();
+        let rm_state = mutex.deref().borrow();
+        Ok(rm_state.fst().borrow().start())
     }
 
     fn compute_trs_and_final_weight(&self, state: StateId) -> Result<(TrsVec<W>, Option<W>)> {
-        let (trs, final_weight) = self.rmeps_state.borrow_mut().expand(state)?;
+        let mut mutex = self.rmeps_state.lock().unwrap();
+        let rm_state = mutex.deref_mut();
+        let (trs, final_weight) = rm_state.expand(state)?;
         let zero = W::zero();
 
         let trs = trs.into_iter().rev().collect_vec();