From be44a705e1affaced72a56e0293006e4cb201b64 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 29 Jul 2024 14:32:44 +0200
Subject: [PATCH 01/13] add dropped_stack to Workspace

---
 src/knuthbendix2.jl       |  1 +
 src/knuthbendix_idxA.jl   | 22 ++++++++++++++++------
 src/settings_workspace.jl |  2 ++
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/knuthbendix2.jl b/src/knuthbendix2.jl
index d0f3b22..51abfc4 100644
--- a/src/knuthbendix2.jl
+++ b/src/knuthbendix2.jl
@@ -112,6 +112,7 @@ function deriverule!(
                 push!(rws, new_rule)
                 deactivate_rules!(rws, stack, new_rule, work)
             else
+                push!(work.dropped_stack, (a, b))
                 work.dropped_rules += 1
             end
         end
diff --git a/src/knuthbendix_idxA.jl b/src/knuthbendix_idxA.jl
index c96110d..2dfa2e0 100644
--- a/src/knuthbendix_idxA.jl
+++ b/src/knuthbendix_idxA.jl
@@ -48,16 +48,27 @@ end
 
 function knuthbendix!(
     settings::Settings{KBIndex},
-    rws::AbstractRewritingSystem{W},
-) where {W}
+    rws::AbstractRewritingSystem,
+)
     if !isreduced(rws)
         rws = reduce!(settings.algorithm, rws)
     end
     # rws is reduced now so we can create its index
     idxA = IndexAutomaton(rws)
     work = Workspace(idxA, settings)
-    stack = Vector{Tuple{W,W}}()
+    knuthbendix!(work, rws, idxA)
+
+    __post!(rws, idxA, work)
+    return rws
+end
 
+function knuthbendix!(
+    work::Workspace{KBIndex},
+    rws::AbstractRewritingSystem{W},
+    idxA::IndexAutomaton = IndexAutomaton(rws),
+) where {W}
+    @assert isreduced(rws)
+    stack = Vector{Tuple{W,W}}()
     rwrules = __rawrules(rws)
     settings = work.settings
 
@@ -133,8 +144,6 @@ function knuthbendix!(
         i += 1
     end
 
-    __post!(rws, idxA, work)
-
     return rws # so the rws is reduced here as well
 end
 
@@ -160,7 +169,7 @@ function __post!(rws::AbstractRewritingSystem, rewriting, work::Workspace)
             if settings.verbosity ≥ 1
                 @warn "The rws does NOT represent the original congruence. Re-adding the missing rules."
             end
-            rws = reduce!(settings.algorithm, rws, stack, work)
+            rws, _ = reduce!(settings.algorithm, rws, stack, 0, 0, work)
         else
             if settings.verbosity == 2
                 @info "Some rules have been dropped but the congruence is preserved."
@@ -169,3 +178,4 @@ function __post!(rws::AbstractRewritingSystem, rewriting, work::Workspace)
     end
     return rws
 end
+
diff --git a/src/settings_workspace.jl b/src/settings_workspace.jl
index 6f56b3d..752f55b 100644
--- a/src/settings_workspace.jl
+++ b/src/settings_workspace.jl
@@ -81,6 +81,7 @@ mutable struct Workspace{CA,T,H,S<:Settings{CA}}
     settings::S
     confluence_timer::Int
     dropped_rules::Int
+    dropped_stack::Vector{Tuple{Word{T},Word{T}}}
 end
 
 function Workspace(word_t, history, settings::Settings)
@@ -91,6 +92,7 @@ function Workspace(word_t, history, settings::Settings)
         settings,
         0,
         0,
+        Tuple{word_t,word_t}[],
     )
 end
 

From fc79cfb7c61ae25176ab366f4c7249899a9bbbc7 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Thu, 11 Jul 2024 12:07:11 +0200
Subject: [PATCH 02/13] add implementation of PrefixAutomaton

---
 Project.toml                     |   4 +-
 src/Automata/Automata.jl         |   2 +
 src/Automata/prefix_automaton.jl | 141 +++++++++++++++++++++++++++++++
 3 files changed, 146 insertions(+), 1 deletion(-)
 create mode 100644 src/Automata/prefix_automaton.jl

diff --git a/Project.toml b/Project.toml
index ac376de..aff9a83 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,16 +7,18 @@ version = "0.4.0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 MacroTools = "0.5"
 ProgressMeter = "1"
+SparseArrays = "1.6"
 julia = "1.6"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "Documenter"]
diff --git a/src/Automata/Automata.jl b/src/Automata/Automata.jl
index c183333..85fb58b 100644
--- a/src/Automata/Automata.jl
+++ b/src/Automata/Automata.jl
@@ -11,6 +11,8 @@ include("interface.jl")
 include("index_automaton.jl")
 include("rebuilding_idxA.jl")
 
+include("prefix_automaton.jl")
+
 include("backtrack.jl")
 
 end # of module Automata
diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
new file mode 100644
index 0000000..38e8928
--- /dev/null
+++ b/src/Automata/prefix_automaton.jl
@@ -0,0 +1,141 @@
+using SparseArrays
+
+struct PrefixAutomaton <: Automaton{Int32}
+    alphabet_len::Int
+    transitions::Vector{SparseVector{Int32,UInt32}}
+    __storage::BitSet
+    # 1 is the initial state
+    # 0 is the fail state
+    # negative values in transitions indicate pointers to
+    # (externally stored) values
+    function PrefixAutomaton(alphabet_len::Integer)
+        transitions = Vector{SparseVector{Int32,UInt32}}(undef, 0)
+        __storage = BitSet()
+        at = new(alphabet_len, transitions, __storage)
+        _ = addstate!(at)
+        return at
+    end
+end
+
+initial(::PrefixAutomaton) = one(Int32)
+
+Base.@propagate_inbounds function hasedge(
+    pfxA::PrefixAutomaton,
+    σ::Integer,
+    label,
+)
+    return pfxA.transitions[σ][label] ≠ 0
+end
+
+Base.@propagate_inbounds function addedge!(
+    pfxA::PrefixAutomaton,
+    src::Integer,
+    dst::Integer,
+    label,
+)
+    pfxA.transitions[src][label] = dst
+    return pfxA
+end
+
+isfail(::PrefixAutomaton, σ::Integer) = iszero(σ)
+isterminal(::PrefixAutomaton, σ::Integer) = σ < 0
+Base.@propagate_inbounds function trace(
+    label::Integer,
+    pfxA::PrefixAutomaton,
+    σ::Integer,
+)
+    return pfxA.transitions[σ][label]
+end
+
+degree(pfxA::PrefixAutomaton, σ::Int32) = nnz(pfxA.transitions[σ])
+
+function addstate!(pfxA::PrefixAutomaton)
+    if !isempty(pfxA.__storage)
+        st = popfirst!(pfxA.__storage)
+        pfxA.transitions[st] .= 0
+        dropzeros!(pfxA.transitions[st])
+        return st
+    else
+        vec = SparseVector{Int32,UInt32}(pfxA.alphabet_len, UInt32[], Int32[])
+        push!(pfxA.transitions, vec)
+        return length(pfxA.transitions)
+    end
+end
+
+function add_direct_path!(pfxA::PrefixAutomaton, rule, val::Integer)
+    @assert val ≤ 0
+    lhs, _ = rule
+    σ = initial(pfxA)
+    for (i, letter) in pairs(lhs)
+        τ = trace(letter, pfxA, σ)
+        # @info "idx = $i" letter τ
+        if i == lastindex(lhs)
+            # if !isfail(pfxA, τ)
+            #     @warn "replacing value $lhs => $σ with" val
+            # end
+            addedge!(pfxA, σ, val, letter)
+            return true, pfxA
+        elseif isfail(pfxA, τ)
+            τ = addstate!(pfxA)
+            addedge!(pfxA, σ, τ, letter)
+        end
+        σ = τ
+        if isterminal(pfxA, σ)
+            @warn "prefix of length $i of $lhs is aready terminal:" σ
+
+            # this may happen if the rule.lhs we push into pfxA
+            # has a prefix that is already in the language of pfxA
+            # then we return false, and we don't enlarge pfxA
+            return false, pfxA
+        end
+    end
+    @error "unintended exit"
+    return false, pfxA
+end
+
+function remove_direct_path!(pfxA::PrefixAutomaton, rule)
+    lhs, _ = rule
+    σ = initial(pfxA)
+    just_before_leaf = σ
+    on_leaf = false
+    k = 0
+
+    for (i, letter) in enumerate(lhs)
+        τ = trace(letter, pfxA, σ)
+        isfail(pfxA, τ) && return pfxA
+        isterminal(pfxA, τ) && i ≠ length(lhs) && return pfxA
+        isterminal(pfxA, τ) && break
+        if degree(pfxA, τ) == 1 && !on_leaf
+            just_before_leaf = σ
+            on_leaf = true
+            k = i
+        end
+        if on_leaf && degree(pfxA, τ) > 1
+            on_leaf = false
+        end
+        σ = τ
+    end
+
+    σ = just_before_leaf
+    for letter in @view lhs[k:end-1]
+        # we're on the "long-leaf" part
+        τ = trace(letter, pfxA, σ)
+        pfxA.transitions[σ][letter] = 0
+        dropzeros!(pfxA.transitions[σ])
+        push!(pfxA.__storage, τ)
+        σ = τ
+    end
+
+    return pfxA
+end
+
+function Base.empty!(pfxA::PrefixAutomaton)
+    union!(pfxA.__storage, 2:length(pfxA.transitions))
+    pfxA.transitions[1] .= 0
+    dropzeros!(pfxA.transitions[1])
+    return pfxA
+end
+
+function Base.isempty(pfxA::PrefixAutomaton)
+    return length(pfxA.transitions) - length(pfxA.__storage) == 1
+end

From 523dedb1195d780cd99e438e49d0f2ab9e940ab4 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 29 Jul 2024 16:21:06 +0200
Subject: [PATCH 03/13] update PrefixAutomaton to isaccepting

---
 src/Automata/prefix_automaton.jl | 38 ++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index 38e8928..d54106d 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -38,7 +38,7 @@ Base.@propagate_inbounds function addedge!(
 end
 
 isfail(::PrefixAutomaton, σ::Integer) = iszero(σ)
-isterminal(::PrefixAutomaton, σ::Integer) = σ < 0
+isaccepting(pfx::PrefixAutomaton, σ::Integer) = 1 ≤ σ ≤ length(pfx.transitions)
 Base.@propagate_inbounds function trace(
     label::Integer,
     pfxA::PrefixAutomaton,
@@ -80,12 +80,12 @@ function add_direct_path!(pfxA::PrefixAutomaton, rule, val::Integer)
             addedge!(pfxA, σ, τ, letter)
         end
         σ = τ
-        if isterminal(pfxA, σ)
-            @warn "prefix of length $i of $lhs is aready terminal:" σ
+        if !isaccepting(pfxA, σ)
+            @debug "prefix of length $i of $lhs is aready a lhs of a rule" σ
 
             # this may happen if the rule.lhs we push into pfxA
-            # has a prefix that is already in the language of pfxA
-            # then we return false, and we don't enlarge pfxA
+            # has a prefix that is reducible; then we return false,
+            # and we don't enlarge pfxA
             return false, pfxA
         end
     end
@@ -96,30 +96,36 @@ end
 function remove_direct_path!(pfxA::PrefixAutomaton, rule)
     lhs, _ = rule
     σ = initial(pfxA)
-    just_before_leaf = σ
     on_leaf = false
-    k = 0
+    leaf_start = (σ, 0)
 
     for (i, letter) in enumerate(lhs)
+        # analyze edge with (src=σ, label=letter, dst=τ)
         τ = trace(letter, pfxA, σ)
         isfail(pfxA, τ) && return pfxA
-        isterminal(pfxA, τ) && i ≠ length(lhs) && return pfxA
-        isterminal(pfxA, τ) && break
-        if degree(pfxA, τ) == 1 && !on_leaf
-            just_before_leaf = σ
-            on_leaf = true
-            k = i
+        if !isaccepting(pfxA, τ)
+            if i == length(lhs)
+                break # we reached the leaf corresponding to lhs
+            end
+            # reached a leaf node before lhs is completed
+            # i.e. lhs does not define a leaf, so there's nothing to remove
+            return pfxA
         end
-        if on_leaf && degree(pfxA, τ) > 1
+        if degree(pfxA, τ) > 1
             on_leaf = false
+        elseif !on_leaf
+            on_leaf = true
+            leaf_start = (σ, i)
         end
         σ = τ
     end
 
-    σ = just_before_leaf
-    for letter in @view lhs[k:end-1]
+    σ, i = leaf_start
+    for letter in @view lhs[i+1:end-1]
         # we're on the "long-leaf" part
         τ = trace(letter, pfxA, σ)
+        # by the early exit above we know there's something to remove
+        @assert isaccepting(pfxA, τ)
         pfxA.transitions[σ][letter] = 0
         dropzeros!(pfxA.transitions[σ])
         push!(pfxA.__storage, τ)

From 73a3c3269cd2cd3d79c734c938e05cc4130985b1 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 29 Jul 2024 16:25:17 +0200
Subject: [PATCH 04/13] pass directly lhs to (add|remove)_direct_path!

---
 src/Automata/prefix_automaton.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index d54106d..1be2aa9 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -62,7 +62,11 @@ function addstate!(pfxA::PrefixAutomaton)
     end
 end
 
-function add_direct_path!(pfxA::PrefixAutomaton, rule, val::Integer)
+function add_direct_path!(
+    pfxA::PrefixAutomaton,
+    lhs::AbstractWord,
+    val::Integer,
+)
     @assert val ≤ 0
     lhs, _ = rule
     σ = initial(pfxA)
@@ -70,9 +74,6 @@ function add_direct_path!(pfxA::PrefixAutomaton, rule, val::Integer)
         τ = trace(letter, pfxA, σ)
         # @info "idx = $i" letter τ
         if i == lastindex(lhs)
-            # if !isfail(pfxA, τ)
-            #     @warn "replacing value $lhs => $σ with" val
-            # end
             addedge!(pfxA, σ, val, letter)
             return true, pfxA
         elseif isfail(pfxA, τ)
@@ -93,8 +94,7 @@ function add_direct_path!(pfxA::PrefixAutomaton, rule, val::Integer)
     return false, pfxA
 end
 
-function remove_direct_path!(pfxA::PrefixAutomaton, rule)
-    lhs, _ = rule
+function remove_direct_path!(pfxA::PrefixAutomaton, lhs::AbstractWord)
     σ = initial(pfxA)
     on_leaf = false
     leaf_start = (σ, 0)

From 4f330101db88651506967f6d550e472999033063 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 29 Jul 2024 16:27:16 +0200
Subject: [PATCH 05/13] don't dropzeros! on transition vectors

the early states will be re-used often and hopefully
will have few structural zeros, while avoiding re-allocation
and memcopy when adding values to sparse vectors
---
 src/Automata/prefix_automaton.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index 1be2aa9..b148adc 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -127,7 +127,6 @@ function remove_direct_path!(pfxA::PrefixAutomaton, lhs::AbstractWord)
         # by the early exit above we know there's something to remove
         @assert isaccepting(pfxA, τ)
         pfxA.transitions[σ][letter] = 0
-        dropzeros!(pfxA.transitions[σ])
         push!(pfxA.__storage, τ)
         σ = τ
     end
@@ -138,7 +137,6 @@ end
 function Base.empty!(pfxA::PrefixAutomaton)
     union!(pfxA.__storage, 2:length(pfxA.transitions))
     pfxA.transitions[1] .= 0
-    dropzeros!(pfxA.transitions[1])
     return pfxA
 end
 

From 1d4c977dd9fc155564103fc2d327971e05d3d229 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Mon, 29 Jul 2024 16:38:41 +0200
Subject: [PATCH 06/13] add PackedVector

---
 src/KnuthBendix.jl         |  1 +
 src/utils/packed_vector.jl | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 src/utils/packed_vector.jl

diff --git a/src/KnuthBendix.jl b/src/KnuthBendix.jl
index 0ec2090..95366db 100644
--- a/src/KnuthBendix.jl
+++ b/src/KnuthBendix.jl
@@ -31,6 +31,7 @@ include("knuthbendix_idxA.jl")
 
 include("confluence_check.jl")
 
+include("utils/packed_vector.jl")
 include("parsing.jl")
 
 include("examples.jl")
diff --git a/src/utils/packed_vector.jl b/src/utils/packed_vector.jl
new file mode 100644
index 0000000..1786b4c
--- /dev/null
+++ b/src/utils/packed_vector.jl
@@ -0,0 +1,36 @@
+
+struct PackedVector{T} <:
+       AbstractVector{SubArray{T,1,Vector{T},Tuple{UnitRange{Int64}},true}}
+    linear_tape::Vector{T}
+    subset_pointers::Vector{Int}
+    PackedVector{T}() where {T} = new{T}(T[], [1])
+end
+
+Base.size(pvec::PackedVector) = (length(pvec.subset_pointers) - 1,)
+Base.@propagate_inbounds function Base.getindex(pvec::PackedVector, i::Integer)
+    @boundscheck 1 ≤ i ≤ length(pvec)
+    ptr = pvec.subset_pointers
+    return @inbounds @view pvec.linear_tape[ptr[i]:ptr[i+1]-1]
+end
+
+__unsafe_push!(pvec::PackedVector, v) = (push!(pvec.linear_tape, v); pvec)
+function __unsafe_finalize!(pvec::PackedVector)
+    k = length(pvec.linear_tape) + 1
+    @assert k ≥ last(pvec.subset_pointers)
+    push!(pvec.subset_pointers, k)
+    return pvec
+end
+
+function Base.resize!(pvec::PackedVector, n::Integer)
+    @assert 0 ≤ n ≤ length(pvec) "growing of pvec is not supported"
+    resize!(pvec.subset_pointers, n + 1)
+    resize!(pvec.linear_tape, pvec.subset_pointers[end] - 1)
+    return pvec
+end
+
+function Base.push!(pvec::PackedVector, v::AbstractVector)
+    append!(pvec.linear_tape, v)
+    push!(pvec.subset_pointers, pvec.subset_pointers[end] + length(v))
+    return pvec
+end
+

From 1b5465a65946ef73e442b4784bc908e3e7f3a248 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Tue, 30 Jul 2024 11:54:33 +0200
Subject: [PATCH 07/13] implement rewrite! based on PrefixAutomaton

---
 src/Automata/Automata.jl         |   2 +-
 src/Automata/prefix_automaton.jl |  42 ++++++---
 src/KnuthBendix.jl               |   2 +-
 src/rewriting.jl                 | 142 +++++++++++++++++++++++++------
 4 files changed, 150 insertions(+), 38 deletions(-)

diff --git a/src/Automata/Automata.jl b/src/Automata/Automata.jl
index 85fb58b..2efb492 100644
--- a/src/Automata/Automata.jl
+++ b/src/Automata/Automata.jl
@@ -4,7 +4,7 @@ import ..KnuthBendix
 import ..KnuthBendix: AbstractWord, RewritingOrdering, RewritingSystem, Rule
 import ..KnuthBendix: alphabet, rules, word_type
 
-export IndexAutomaton
+export IndexAutomaton, PrefixAutomaton
 
 include("states.jl")
 include("interface.jl")
diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index b148adc..5612d35 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -1,19 +1,26 @@
 using SparseArrays
 
-struct PrefixAutomaton <: Automaton{Int32}
+struct PrefixAutomaton{V} <: Automaton{Int32}
     alphabet_len::Int
     transitions::Vector{SparseVector{Int32,UInt32}}
     __storage::BitSet
+    rwrules::V
     # 1 is the initial state
     # 0 is the fail state
     # negative values in transitions indicate pointers to
     # (externally stored) values
-    function PrefixAutomaton(alphabet_len::Integer)
+    function PrefixAutomaton(
+        alphabet_len::Integer,
+        rules::V,
+    ) where {V<:AbstractVector}
         transitions = Vector{SparseVector{Int32,UInt32}}(undef, 0)
         __storage = BitSet()
-        at = new(alphabet_len, transitions, __storage)
-        _ = addstate!(at)
-        return at
+        pfxA = new{V}(alphabet_len, transitions, __storage, rules)
+        _ = addstate!(pfxA)
+        for (i, rule) in pairs(rules)
+            add_direct_path!(pfxA, rule.lhs, -i)
+        end
+        return pfxA
     end
 end
 
@@ -39,11 +46,8 @@ end
 
 isfail(::PrefixAutomaton, σ::Integer) = iszero(σ)
 isaccepting(pfx::PrefixAutomaton, σ::Integer) = 1 ≤ σ ≤ length(pfx.transitions)
-Base.@propagate_inbounds function trace(
-    label::Integer,
-    pfxA::PrefixAutomaton,
-    σ::Integer,
-)
+
+@inline function trace(label::Integer, pfxA::PrefixAutomaton, σ::Integer)
     return pfxA.transitions[σ][label]
 end
 
@@ -68,7 +72,6 @@ function add_direct_path!(
     val::Integer,
 )
     @assert val ≤ 0
-    lhs, _ = rule
     σ = initial(pfxA)
     for (i, letter) in pairs(lhs)
         τ = trace(letter, pfxA, σ)
@@ -143,3 +146,20 @@ end
 function Base.isempty(pfxA::PrefixAutomaton)
     return length(pfxA.transitions) - length(pfxA.__storage) == 1
 end
+
+function PrefixAutomaton(rws::AbstractRewritingSystem)
+    l = length(alphabet(rws))
+    return PrefixAutomaton(l, deepcopy(KnuthBendix.__rawrules(rws)))
+end
+
+function Base.show(io::IO, ::MIME"text/plain", pfxA::PrefixAutomaton)
+    println(
+        io,
+        "prefix automaton over alphabet of $(pfxA.alphabet_len) letters:",
+    )
+    println(io, "  • $(length(pfxA.transitions)) accepting states")
+    nrules = mapreduce(+, pairs(pfxA.transitions)) do (i, t)
+        return i in pfxA.__storage ? 0 : sum(<(0), t)
+    end
+    return print(io, "  • $(nrules) non-accepting states (rw rules)")
+end
diff --git a/src/KnuthBendix.jl b/src/KnuthBendix.jl
index 95366db..f6678de 100644
--- a/src/KnuthBendix.jl
+++ b/src/KnuthBendix.jl
@@ -6,6 +6,7 @@ export Alphabet, Word, RewritingSystem
 export LenLex, WreathOrder, Recursive, WeightedLex
 export alphabet, isconfluent, ordering, knuthbendix
 
+include("utils/packed_vector.jl")
 include("Words/Words.jl")
 using .Words
 include("buffer_pair.jl")
@@ -31,7 +32,6 @@ include("knuthbendix_idxA.jl")
 
 include("confluence_check.jl")
 
-include("utils/packed_vector.jl")
 include("parsing.jl")
 
 include("examples.jl")
diff --git a/src/rewriting.jl b/src/rewriting.jl
index b1c9bb1..bda4420 100644
--- a/src/rewriting.jl
+++ b/src/rewriting.jl
@@ -1,5 +1,8 @@
 RewritingBuffer{T}(::Any) where {T} = RewritingBuffer{T}() # no history
 RewritingBuffer{T}(::IndexAutomaton{S}) where {T,S} = RewritingBuffer{T}(S[])
+function RewritingBuffer{T}(::PrefixAutomaton) where {T,S}
+    return RewritingBuffer{T}(PackedVector{UInt32}())
+end
 
 """
     rewrite(u::AbstractWord, rewriting)
@@ -39,6 +42,42 @@ function rewrite(
     return W(v)
 end
 
+"""
+    function rewrite!(bp::BufferPair, rewriting; kwargs...)
+Rewrites word stored in `BufferPair` using `rewriting` object.
+
+To store a word in `bp`
+[`Words.store!`](@ref Words.store!(::BufferPair, ::AbstractWord))
+should be used.
+
+!!! warning
+    This implementation returns an instance of `Words.BufferWord` aliased with
+    the intenrals of `BufferPair`. You need to copy the return value if you
+    want to take the ownership.
+"""
+function rewrite!(rwb::RewritingBuffer, rewriting; kwargs...)
+    v = if isempty(rewriting)
+        Words.store!(rwb.output, rwb.input)
+    else
+        rewrite!(rwb.output, rwb.input, rewriting; kwargs...)
+    end
+    empty!(rwb.input) # shifts bp.input pointers to the beginning of its storage
+    return v
+end
+function rewrite!(
+    rwb::RewritingBuffer,
+    rewriting::Automata.Automaton;
+    kwargs...,
+)
+    v = if isempty(rewriting)
+        Words.store!(rwb.output, rwb.input)
+    else
+        rewrite!(rwb.output, rwb.input, rewriting; history = rwb.history, kwargs...)
+    end
+    empty!(rwb.input) # shifts bp.input pointers to the beginning of its storage
+    return v
+end
+
 function rewrite!(v::AbstractWord, w::AbstractWord, A::Any; kwargs...)
     throw(
         """No method for rewriting with $(typeof(A)). You need to implement
@@ -182,7 +221,7 @@ Whenever a terminal (i.e. accepting) state is encountered
 2. the appropriate suffix of `v` (equal to `lhs`) is removed, and
 3. `rhs` is prepended to `w`.
 
-Tracing continues from the newly prepended letter.
+Tracing continues from the first letter of the newly prepended word.
 
 To continue tracing `w` through the automaton we need to backtrack on our path
 in the automaton and for this `rewrite` maintains a vector of visited states of
@@ -233,33 +272,86 @@ function rewrite!(
 end
 
 """
-    function rewrite!(bp::BufferPair, rewriting; kwargs...)
-Rewrites word stored in `BufferPair` using `rewriting` object.
+    rewrite!(v::AbstractWord, w::AbstractWord, pfxA::PrefixAutomaton[; history, skipping])
+Rewrite word `w` storing the result in `v` using prefix automaton `idxA`.
+As rewriting rules are stored **externally**, they must be passed in the
+`rules` keyword argument.
 
-To store a word in `bp`
-[`Words.store!`](@ref Words.store!(::BufferPair, ::AbstractWord))
-should be used.
+Rewriting with a [`PrefixAutomaton`](@ref Automata.PrefixAutomaton) traces
+(i.e. follows) simultanously all paths in the automaton determined by `w`.
+To be more precise we trace a path in the power-set automaton (states are the
+subsets of states of the original automaton) via lazy accessible set construction.
+Since the non-deterministic part of the automaton consists of `ε`-loop at the
+initial state there are at most `length(w)-1` such paths.
 
-!!! warning
-    This implementation returns an instance of `Words.BufferWord` aliased with
-    the intenrals of `BufferPair`. You need to copy the return value if you
-    want to take the ownership.
+Whenever a non-accepting state is encountered **on any** of those paths
+
+1. its corresponding rule `lhs → rhs` is retrived,
+2. the appropriate suffix of `v` (equal to `lhs`) is removed, and
+3. `rhs` is prepended to `w`.
+
+Tracing continues from the first letter of the newly prepended word.
+
+To continue tracing `w` through the automaton we need to backtrack on our path
+in the automaton and for this `rewrite` maintains a history of visited
+states of `pfxA`. Whenever a suffix is removed from `v`, the path is rewinded
+(i.e. shortened to the appropriate length) and the next letter of `w` is traced
+from the last state on the path. This maintains the property that signature of
+the path is equal to `v` at all times.
+
+Once prefix automaton is build the complexity of this rewriting is `Ω(length(w)²)`.
 """
-function rewrite!(bp::RewritingBuffer, rewriting; kwargs...)
-    v = if isempty(rewriting)
-        Words.store!(bp.output, bp.input)
-    else
-        rewrite!(bp.output, bp.input, rewriting; kwargs...)
-    end
-    empty!(bp.input) # shifts bp._wWord pointers to the beginning of its storage
-    return v
-end
-function rewrite!(bp::RewritingBuffer, rewriting::IndexAutomaton; kwargs...)
-    v = if isempty(rewriting)
-        Words.store!(bp.output, bp.input)
-    else
-        rewrite!(bp.output, bp.input, rewriting; history = bp.history, kwargs...)
+function rewrite!(
+    v::AbstractWord,
+    w::AbstractWord,
+    pfxA::PrefixAutomaton;
+    history::PackedVector = PackedVector{UInt32}(),
+    skipping = nothing,
+)
+    resize!(history, 0)
+    __unsafe_push!(history, Automata.initial(pfxA))
+    __unsafe_finalize!(history)
+    v = resize!(v, 0)
+    # we're doing path tracing on PrefixAutomaton that is non-deterministic
+    # in the sense that we add an ε-loop at the initial state
+    # thus this is path tracing via (lazy) accessible set construction
+    # i.e. simultanously tracing all possible paths with the given signature,
+    # or tracing a path in power-set automaton (states are the subsets of
+    # states of the original automaton).
+    # We rewind the history of ALL paths whenever a terminal is found in ONE of them.
+    while !isone(w)
+        letter = popfirst!(w)
+        # we're tracing a bunch of paths simultanously:
+        found_terminal = false
+        # @info "current multi-state" last(history)
+        for σ in last(history)
+            τ = Automata.trace(letter, pfxA, σ)
+            # @info "with letter=$letter we transition" src = σ dst = τ
+            Automata.isfail(pfxA, τ) && continue # this path doesn't proceed any further
+            if !Automata.isaccepting(pfxA, τ) && -τ ≠ skipping
+                # find the length of the corresponding lhs and rewind
+                # @info "The dst is terminal, using:" pfxA.rwrules[-τ]
+                lhs, rhs = pfxA.rwrules[-τ]
+                resize!(v, length(v) - length(lhs) + 1)
+                prepend!(w, rhs)
+                resize!(history, length(history) - length(lhs) + 1)
+                found_terminal = true
+                break
+            end
+            if !found_terminal
+                __unsafe_push!(history, τ)
+            end
+        end
+        if !found_terminal
+            # @info """none of the dsts were terminal:
+            # extending v (by $letter) & pushing initial (1) to history"""
+            push!(v, letter)
+            # we finish by a suffix of w by adding the initial state:
+            __unsafe_push!(history, Automata.initial(pfxA))
+            # after we're done with all of the path we proclaim the next subset
+            __unsafe_finalize!(history)
+        end
+        # @info "afterwards:" v w
     end
-    empty!(bp.input) # shifts bp._wWord pointers to the beginning of its storage
     return v
 end

From 17bf038078998e0d9edc8d012c37bcb8c512ba38 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Tue, 30 Jul 2024 11:55:36 +0200
Subject: [PATCH 08/13] import AbstractRewritingSystem to Automata

---
 src/Automata/Automata.jl        | 3 ++-
 src/Automata/index_automaton.jl | 2 +-
 src/Automata/rebuilding_idxA.jl | 4 ++--
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/Automata/Automata.jl b/src/Automata/Automata.jl
index 2efb492..246a60b 100644
--- a/src/Automata/Automata.jl
+++ b/src/Automata/Automata.jl
@@ -1,7 +1,8 @@
 module Automata
 
 import ..KnuthBendix
-import ..KnuthBendix: AbstractWord, RewritingOrdering, RewritingSystem, Rule
+import ..KnuthBendix:
+    AbstractWord, RewritingOrdering, AbstractRewritingSystem, Rule
 import ..KnuthBendix: alphabet, rules, word_type
 
 export IndexAutomaton, PrefixAutomaton
diff --git a/src/Automata/index_automaton.jl b/src/Automata/index_automaton.jl
index fd84c17..8bfee92 100644
--- a/src/Automata/index_automaton.jl
+++ b/src/Automata/index_automaton.jl
@@ -51,7 +51,7 @@ Base.Base.@propagate_inbounds function trace(
     return σ[label]
 end
 
-function IndexAutomaton(rws::RewritingSystem{W}) where {W}
+function IndexAutomaton(rws::AbstractRewritingSystem{W}) where {W}
     if !KnuthBendix.isreduced(rws)
         throw(
             ArgumentError(
diff --git a/src/Automata/rebuilding_idxA.jl b/src/Automata/rebuilding_idxA.jl
index 7d9b6e6..51c1925 100644
--- a/src/Automata/rebuilding_idxA.jl
+++ b/src/Automata/rebuilding_idxA.jl
@@ -1,4 +1,4 @@
-function _rebuild!(idxA::IndexAutomaton, rws::RewritingSystem)
+function _rebuild!(idxA::IndexAutomaton, rws::AbstractRewritingSystem)
     # Most of the information in idxA can be reused;
     # however here we just rebuild it from scratch
     at = IndexAutomaton(rws)
@@ -8,7 +8,7 @@ function _rebuild!(idxA::IndexAutomaton, rws::RewritingSystem)
     return idxA
 end
 
-function rebuild!(idxA::IndexAutomaton, rws::RewritingSystem)
+function rebuild!(idxA::IndexAutomaton, rws::AbstractRewritingSystem)
     # mark all states as not up to date
     for states in idxA.states
         for σ in states

From 908465fd84c63a0819e28f2d0ccbd26e7efd150f Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Tue, 30 Jul 2024 11:56:01 +0200
Subject: [PATCH 09/13] fix: general tracing through Automata

---
 src/Automata/interface.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/Automata/interface.jl b/src/Automata/interface.jl
index a2acb99..c1d8c6d 100644
--- a/src/Automata/interface.jl
+++ b/src/Automata/interface.jl
@@ -55,11 +55,11 @@ returned.
     for (i, l) in enumerate(w)
         if hasedge(A, σ, l)
             τ = trace(l, A, σ)
+            isfail(A, τ) && return i - 1, σ
+            σ = τ
+        else
+            return return i - 1, σ
         end
-        if isfail(A, τ)
-            return i - 1, σ
-        end
-        σ = τ
     end
     return length(w), σ
 end

From d264049c025751c75276fae0f7585d64020e55b2 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Tue, 30 Jul 2024 12:29:11 +0200
Subject: [PATCH 10/13] add tests for PackedVector and rewriting with automata

---
 test/automata.jl      | 11 +++++++--
 test/kbs.jl           | 53 +++++++++++++++++++++++++++++++++++++++++--
 test/kbs1.jl          |  2 +-
 test/packed_vector.jl | 30 ++++++++++++++++++++++++
 test/runtests.jl      |  1 +
 5 files changed, 92 insertions(+), 5 deletions(-)
 create mode 100644 test/packed_vector.jl

diff --git a/test/automata.jl b/test/automata.jl
index c71cbeb..3661120 100644
--- a/test/automata.jl
+++ b/test/automata.jl
@@ -69,8 +69,10 @@ end
     rs =
         KB.RewritingSystem([(A, ε), (B, ε)], lenlexord, reduced = true)
     ia = KB.IndexAutomaton(rs)
+    pa = KB.PrefixAutomaton(rs)
     @test KB.rewrite(testword, rs) ==
           KB.rewrite(testword, ia) ==
+          KB.rewrite(testword, pa) ==
           Word([1])
 
     rs = KB.RewritingSystem(
@@ -88,13 +90,18 @@ end
         reduced = true,
     )
     ia = KB.IndexAutomaton(rs)
+    pa = KB.PrefixAutomaton(rs)
 
     @test !isempty(ia)
+    @test !isempty(pa)
 
-    @test KB.rewrite(testword, rs) == KB.rewrite(testword, ia)
+    @test KB.rewrite(testword, rs) ==
+          KB.rewrite(testword, ia) ==
+          KB.rewrite(testword, pa)
 
     w = Word([1, 3, 4, 1, 4, 4, 1, 1, 4, 2, 3, 2, 4, 2, 2, 3, 1, 2, 1])
-    @test KB.rewrite(w, rs) == KB.rewrite(w, ia)
+    @test KB.rewrite(w, rs) == KB.rewrite(w, ia) == KB.rewrite(w, pa)
 
     @test sprint(show, ia) isa String
+    @test sprint(show, pa) isa String
 end
diff --git a/test/kbs.jl b/test/kbs.jl
index 859c498..2b9afdb 100644
--- a/test/kbs.jl
+++ b/test/kbs.jl
@@ -1,4 +1,4 @@
-@testset "KBS" begin
+@testset "KBStack" begin
     lenlex = let A = Alphabet([:a, :b, :A, :B])
         KB.setinverse!(A, :a, :A)
         KB.setinverse!(A, :b, :B)
@@ -19,7 +19,6 @@
     @test Set(KB.rules(knuthbendix(KB.Settings(KB.KBPlain()), R))) == crs
     @test Set(KB.rules(knuthbendix(KB.Settings(KB.KBStack()), R))) == crs
     @test Set(KB.rules(knuthbendix(KB.Settings(KB.KBS2AlgRuleDel()), R))) == crs
-    @test Set(KB.rules(knuthbendix(KB.Settings(KB.KBIndex()), R))) == crs
 
     @testset "io for RewritingSystem" begin
         @test sprint(show, MIME"text/plain"(), RC) isa String
@@ -37,3 +36,53 @@
         @test occursin("• verbosity            : 2", res)
     end
 end
+
+@testset "Automaton rewriting" begin
+    rws = KB.ExampleRWS.Hurwitz4()
+    RC = knuthbendix(KB.Settings(KB.KBStack()), rws)
+
+    idxA = KB.IndexAutomaton(RC)
+    pfxA = KB.PrefixAutomaton(RC)
+
+    w = Word([1, 1, 1, 2, 2, 2, 1, 1, 3, 2, 3, 3, 1, 1, 1, 3, 3, 1, 2, 2, 3, 2])
+    v = Word([1, 2, 1, 2, 1, 3])
+    @test v == KB.rewrite(w, RC)
+    @test v == KB.rewrite(w, idxA)
+    @test v == KB.rewrite(w, pfxA)
+
+    let at = idxA
+        rwb = KB.RewritingBuffer{UInt16}(at)
+        k = @allocated begin
+            KB.Words.store!(rwb, w)
+            KB.rewrite!(rwb, at)
+        end
+        @test length(w) == 22
+        @test rwb.output == v
+
+        @test 0 == @allocated begin
+            KB.Words.store!(rwb, w)
+            KB.rewrite!(rwb, at)
+        end
+    end
+
+    let at = pfxA
+        rwb = KB.RewritingBuffer{UInt16}(at)
+        k = @allocated begin
+            KB.Words.store!(rwb, w)
+            KB.rewrite!(rwb, at)
+        end
+        @test length(w) == 22
+        @test rwb.output == v
+
+        @test 0 == @allocated begin
+            KB.Words.store!(rwb, w)
+            KB.rewrite!(rwb, at)
+        end
+    end
+
+    @test all(1:10) do _
+        w = Word(rand(1:length(alphabet(rws)), 50))
+        v = KB.rewrite(w, RC)
+        return v == KB.rewrite(w, idxA) == KB.rewrite(w, pfxA)
+    end
+end
diff --git a/test/kbs1.jl b/test/kbs1.jl
index 3f1e822..b1c832a 100644
--- a/test/kbs1.jl
+++ b/test/kbs1.jl
@@ -1,4 +1,4 @@
-@testset "KBS1" begin
+@testset "KBPlain" begin
     Al = Alphabet([:a, :A, :b, :B])
     KB.setinverse!(Al, :a, :A)
     KB.setinverse!(Al, :b, :B)
diff --git a/test/packed_vector.jl b/test/packed_vector.jl
new file mode 100644
index 0000000..69a2bf3
--- /dev/null
+++ b/test/packed_vector.jl
@@ -0,0 +1,30 @@
+@testset "PackedVector" begin
+    pvec = KB.PackedVector{Char}()
+    @test pvec isa AbstractVector
+    @test length(pvec) == 0
+
+    push!(pvec, 'a':'c')
+    @test length(pvec) == 1
+    @test pvec == [['a', 'b', 'c']]
+
+    push!(pvec, 'd':'g')
+    @test length(pvec) == 2
+    @test pvec == [['a', 'b', 'c'], ['d', 'e', 'f', 'g']]
+
+    KB.__unsafe_finalize!(pvec)
+    @test length(pvec) == 3
+    @test pvec == [['a', 'b', 'c'], ['d', 'e', 'f', 'g'], Char[]]
+
+    KB.__unsafe_push!(pvec, 'a')
+    KB.__unsafe_push!(pvec, 'a')
+    @test length(pvec) == 3
+    @test pvec == [['a', 'b', 'c'], ['d', 'e', 'f', 'g'], Char[]]
+
+    KB.__unsafe_finalize!(pvec)
+    @test length(pvec) == 4
+    @test pvec == [['a', 'b', 'c'], ['d', 'e', 'f', 'g'], Char[], ['a', 'a']]
+
+    @test_throws AssertionError resize!(pvec, 5)
+    @test resize!(pvec, 2) == [['a', 'b', 'c'], ['d', 'e', 'f', 'g']]
+    @test length(pvec) == 2
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 59bc25b..b71549c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,6 +7,7 @@ import KnuthBendix.Automata
 include("abstract_words.jl")
 
 @testset "KnuthBendix.jl" begin
+    include("packed_vector.jl")
     include("words.jl")
     include("bufferwords.jl")
     include("alphabets.jl")

From 7073dcbebe7ae79d04e9774a65e8c53a0823e426 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Wed, 31 Jul 2024 09:53:05 +0200
Subject: [PATCH 11/13] store ordering in PrefixAutomaton

---
 src/Automata/Automata.jl         |  2 +-
 src/Automata/index_automaton.jl  |  2 +-
 src/Automata/prefix_automaton.jl | 12 ++++++------
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/Automata/Automata.jl b/src/Automata/Automata.jl
index 246a60b..0b179a2 100644
--- a/src/Automata/Automata.jl
+++ b/src/Automata/Automata.jl
@@ -3,7 +3,7 @@ module Automata
 import ..KnuthBendix
 import ..KnuthBendix:
     AbstractWord, RewritingOrdering, AbstractRewritingSystem, Rule
-import ..KnuthBendix: alphabet, rules, word_type
+import ..KnuthBendix: alphabet, ordering, rules, word_type
 
 export IndexAutomaton, PrefixAutomaton
 
diff --git a/src/Automata/index_automaton.jl b/src/Automata/index_automaton.jl
index 8bfee92..615e325 100644
--- a/src/Automata/index_automaton.jl
+++ b/src/Automata/index_automaton.jl
@@ -179,7 +179,7 @@ function Base.show(io::IO, idxA::IndexAutomaton)
         count(st -> !Automata.isaccepting(idxA, st), states) for
         states in idxA.states
     ]
-    ord = KnuthBendix.ordering(idxA)
+    ord = ordering(idxA)
     A = alphabet(ord)
     println(io, "index automaton over $(typeof(ord)) with $(length(A)) letters")
     nstates = sum(length, idxA.states)
diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index 5612d35..40c65cd 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -1,21 +1,21 @@
 using SparseArrays
 
-struct PrefixAutomaton{V} <: Automaton{Int32}
-    alphabet_len::Int
+struct PrefixAutomaton{O<:RewritingOrdering,V} <: Automaton{Int32}
+    ordering::O
     transitions::Vector{SparseVector{Int32,UInt32}}
     __storage::BitSet
     rwrules::V
     # 1 is the initial state
     # 0 is the fail state
-    # negative values in transitions indicate pointers to
-    # (externally stored) values
+    # negative values in transitions indicate indices to values stored in rwrules
+
     function PrefixAutomaton(
-        alphabet_len::Integer,
+        ordering::RewritingOrdering,
         rules::V,
     ) where {V<:AbstractVector}
         transitions = Vector{SparseVector{Int32,UInt32}}(undef, 0)
         __storage = BitSet()
-        pfxA = new{V}(alphabet_len, transitions, __storage, rules)
+        pfxA = new{typeof(ordering),V}(ordering, transitions, __storage, rules)
         _ = addstate!(pfxA)
         for (i, rule) in pairs(rules)
             add_direct_path!(pfxA, rule.lhs, -i)

From 2e194aa14d6b604ceb9762aadb76b96b4d35fc9a Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Wed, 31 Jul 2024 09:53:55 +0200
Subject: [PATCH 12/13] re-order functions in (index|prefix)_automaton.jl

---
 src/Automata/index_automaton.jl  | 82 +++++++++++++++++---------------
 src/Automata/prefix_automaton.jl | 76 ++++++++++++++++-------------
 2 files changed, 88 insertions(+), 70 deletions(-)

diff --git a/src/Automata/index_automaton.jl b/src/Automata/index_automaton.jl
index 615e325..d79cda2 100644
--- a/src/Automata/index_automaton.jl
+++ b/src/Automata/index_automaton.jl
@@ -26,24 +26,17 @@ mutable struct IndexAutomaton{S,O<:RewritingOrdering} <: Automaton{S}
 end
 
 initial(idxA::IndexAutomaton) = idxA.initial
-KnuthBendix.ordering(idxA::IndexAutomaton) = idxA.ordering
-
-hasedge(::IndexAutomaton, ::State, ::Integer) = true
-
-addedge!(idxA::IndexAutomaton, src::State, dst::State, label) = src[label] = dst
-
 isfail(idxA::IndexAutomaton, σ::State) = σ === idxA.fail
-isaccepting(idxA::IndexAutomaton, σ::State) = !isdefined(σ, :value)
+isaccepting(::IndexAutomaton, σ::State) = !isdefined(σ, :value)
 
-signature(idxA::IndexAutomaton, σ::State) = id(σ)
-
-Base.isempty(idxA::IndexAutomaton) = degree(initial(idxA)) == 0
+hasedge(::IndexAutomaton, ::State, ::Integer) = true
 
-function KnuthBendix.word_type(::Type{<:IndexAutomaton{S}}) where {S}
-    return eltype(valtype(S))
+function addedge!(idxA::IndexAutomaton, src::State, dst::State, label)
+    src[label] = dst
+    return idxA
 end
 
-Base.Base.@propagate_inbounds function trace(
+function trace(
     label::Integer,
     ::IndexAutomaton,
     σ::State,
@@ -51,32 +44,11 @@ Base.Base.@propagate_inbounds function trace(
     return σ[label]
 end
 
-function IndexAutomaton(rws::AbstractRewritingSystem{W}) where {W}
-    if !KnuthBendix.isreduced(rws)
-        throw(
-            ArgumentError(
-                """`IndexAutomaton` can be constructed from reduced rewriting systems only.
-                Call `KnuthBendix.reduce!(rws)` and try again.""",
-            ),
-        )
-    end
-
-    id = @view one(W)[1:0]
-    S = State{typeof(id),UInt32,eltype(rules(rws))}
-    ord = KnuthBendix.ordering(rws)
-    A = alphabet(ord)
-    fail = S(Vector{S}(undef, length(A)), id, 0)
-    α = State(fail, id, 0)
+Base.isempty(idxA::IndexAutomaton) = degree(initial(idxA)) == 0
 
-    idxA = IndexAutomaton(ord, α, fail, Vector{typeof(α)}[])
-    idxA = self_complete!(idxA, fail, override = true)
-    idxA = direct_edges!(idxA, rules(rws))
-    idxA = skew_edges!(idxA)
+signature(::IndexAutomaton, σ::State) = id(σ)
 
-    return idxA
-end
-
-KnuthBendix.isreduced(idxA::Automata.IndexAutomaton) = true
+# construction/modification
 
 function direct_edges!(idxA::IndexAutomaton, rwrules)
     for (idx, rule) in enumerate(rwrules)
@@ -182,8 +154,42 @@ function Base.show(io::IO, idxA::IndexAutomaton)
     ord = ordering(idxA)
     A = alphabet(ord)
     println(io, "index automaton over $(typeof(ord)) with $(length(A)) letters")
-    nstates = sum(length, idxA.states)
+    nstates = sum(length, idxA.states) + 1 # the initial one
     println(io, "  • ", nstates, " state" * (nstates == 1 ? "" : "s"))
     print(io, "  • ", sum(rules_count), " non-accepting states (rw rules)")
     return
 end
+
+# for using IndexAutomaton as rewriting struct in KnuthBendix
+KnuthBendix.ordering(idxA::IndexAutomaton) = idxA.ordering
+
+function KnuthBendix.word_type(::Type{<:IndexAutomaton{S}}) where {S}
+    return eltype(valtype(S))
+end
+
+function IndexAutomaton(rws::AbstractRewritingSystem{W}) where {W}
+    if !KnuthBendix.isreduced(rws)
+        throw(
+            ArgumentError(
+                """`IndexAutomaton` can be constructed from reduced rewriting systems only.
+                Call `KnuthBendix.reduce!(rws)` and try again.""",
+            ),
+        )
+    end
+
+    id = @view one(W)[1:0]
+    S = State{typeof(id),UInt32,eltype(rules(rws))}
+    ord = ordering(rws)
+    A = alphabet(ord)
+    fail = S(Vector{S}(undef, length(A)), id, 0)
+    α = State(fail, id, 0)
+
+    idxA = IndexAutomaton(ord, α, fail, Vector{typeof(α)}[])
+    idxA = self_complete!(idxA, fail, override = true)
+    idxA = direct_edges!(idxA, rules(rws))
+    idxA = skew_edges!(idxA)
+
+    return idxA
+end
+
+KnuthBendix.isreduced(::IndexAutomaton) = true
diff --git a/src/Automata/prefix_automaton.jl b/src/Automata/prefix_automaton.jl
index 40c65cd..c02fbaf 100644
--- a/src/Automata/prefix_automaton.jl
+++ b/src/Automata/prefix_automaton.jl
@@ -25,16 +25,12 @@ struct PrefixAutomaton{O<:RewritingOrdering,V} <: Automaton{Int32}
 end
 
 initial(::PrefixAutomaton) = one(Int32)
+isfail(::PrefixAutomaton, σ::Integer) = iszero(σ)
+isaccepting(pfx::PrefixAutomaton, σ::Integer) = 1 ≤ σ ≤ length(pfx.transitions)
 
-Base.@propagate_inbounds function hasedge(
-    pfxA::PrefixAutomaton,
-    σ::Integer,
-    label,
-)
-    return pfxA.transitions[σ][label] ≠ 0
-end
+hasedge(pfxA::PrefixAutomaton, σ::Integer, lab) = pfxA.transitions[σ][lab] ≠ 0
 
-Base.@propagate_inbounds function addedge!(
+function addedge!(
     pfxA::PrefixAutomaton,
     src::Integer,
     dst::Integer,
@@ -44,23 +40,31 @@ Base.@propagate_inbounds function addedge!(
     return pfxA
 end
 
-isfail(::PrefixAutomaton, σ::Integer) = iszero(σ)
-isaccepting(pfx::PrefixAutomaton, σ::Integer) = 1 ≤ σ ≤ length(pfx.transitions)
-
-@inline function trace(label::Integer, pfxA::PrefixAutomaton, σ::Integer)
+function trace(label::Integer, pfxA::PrefixAutomaton, σ::Integer)
     return pfxA.transitions[σ][label]
 end
 
-degree(pfxA::PrefixAutomaton, σ::Int32) = nnz(pfxA.transitions[σ])
+function Base.isempty(pfxA::PrefixAutomaton)
+    return all(iszero, pfxA.transitions[initial(pfxA)])
+end
+
+function Base.empty!(pfxA::PrefixAutomaton)
+    union!(pfxA.__storage, 2:length(pfxA.transitions))
+    pfxA.transitions[1] .= 0
+    return pfxA
+end
+
+# construction/modification
 
 function addstate!(pfxA::PrefixAutomaton)
     if !isempty(pfxA.__storage)
         st = popfirst!(pfxA.__storage)
         pfxA.transitions[st] .= 0
-        dropzeros!(pfxA.transitions[st])
+        # dropzeros!(pfxA.transitions[st])
         return st
     else
-        vec = SparseVector{Int32,UInt32}(pfxA.alphabet_len, UInt32[], Int32[])
+        l = length(alphabet(ordering(pfxA)))
+        vec = SparseVector(l, UInt32[], Int32[])
         push!(pfxA.transitions, vec)
         return length(pfxA.transitions)
     end
@@ -137,29 +141,37 @@ function remove_direct_path!(pfxA::PrefixAutomaton, lhs::AbstractWord)
     return pfxA
 end
 
-function Base.empty!(pfxA::PrefixAutomaton)
-    union!(pfxA.__storage, 2:length(pfxA.transitions))
-    pfxA.transitions[1] .= 0
-    return pfxA
-end
-
-function Base.isempty(pfxA::PrefixAutomaton)
-    return length(pfxA.transitions) - length(pfxA.__storage) == 1
-end
-
-function PrefixAutomaton(rws::AbstractRewritingSystem)
-    l = length(alphabet(rws))
-    return PrefixAutomaton(l, deepcopy(KnuthBendix.__rawrules(rws)))
-end
-
 function Base.show(io::IO, ::MIME"text/plain", pfxA::PrefixAutomaton)
+    ord = ordering(pfxA)
+    A = alphabet(ord)
     println(
         io,
-        "prefix automaton over alphabet of $(pfxA.alphabet_len) letters:",
+        "prefix automaton over $(typeof(ord)) with $(length(A)) letters",
     )
-    println(io, "  • $(length(pfxA.transitions)) accepting states")
+    accept_states = length(pfxA.transitions) - length(pfxA.__storage)
     nrules = mapreduce(+, pairs(pfxA.transitions)) do (i, t)
         return i in pfxA.__storage ? 0 : sum(<(0), t)
     end
+    println(io, "  • $(accept_states+nrules) states")
     return print(io, "  • $(nrules) non-accepting states (rw rules)")
 end
+
+function Base.push!(pfxA::PrefixAutomaton, rule::KnuthBendix.Rule)
+    n = length(pfxA.rwrules) + 1
+    added, pfxA = add_direct_path!(pfxA, rule.lhs, -n)
+    if added
+        push!(pfxA.rwrules, rule)
+    end
+    return pfxA
+end
+
+# for using IndexAutomaton as rewriting struct in KnuthBendix
+KnuthBendix.ordering(pfxA::PrefixAutomaton) = pfxA.ordering
+
+function KnuthBendix.word_type(::Type{<:PrefixAutomaton{O,V}}) where {O,V}
+    return KnuthBendix.word_type(eltype(V))
+end
+
+function PrefixAutomaton(rws::AbstractRewritingSystem)
+    return PrefixAutomaton(ordering(rws), KnuthBendix.__rawrules(rws))
+end

From 64691e6138dadff0fe345ef7041c8ea9850c8ec1 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Wed, 31 Jul 2024 09:57:27 +0200
Subject: [PATCH 13/13] define Stack{W} as alias to Vector{Tuple{W,W}}

---
 src/KnuthBendix.jl        | 2 ++
 src/confluence_check.jl   | 2 +-
 src/knuthbendix2.jl       | 2 +-
 src/knuthbendix_base.jl   | 3 ++-
 src/knuthbendix_delete.jl | 2 +-
 src/knuthbendix_idxA.jl   | 5 +++--
 src/rewriting_system.jl   | 2 +-
 src/rules.jl              | 2 ++
 src/settings_workspace.jl | 4 ++--
 9 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/KnuthBendix.jl b/src/KnuthBendix.jl
index f6678de..ab2d7e0 100644
--- a/src/KnuthBendix.jl
+++ b/src/KnuthBendix.jl
@@ -6,6 +6,8 @@ export Alphabet, Word, RewritingSystem
 export LenLex, WreathOrder, Recursive, WeightedLex
 export alphabet, isconfluent, ordering, knuthbendix
 
+const Stack{W} = Vector{Tuple{W,W}}
+
 include("utils/packed_vector.jl")
 include("Words/Words.jl")
 using .Words
diff --git a/src/confluence_check.jl b/src/confluence_check.jl
index e144ac6..3d3230f 100644
--- a/src/confluence_check.jl
+++ b/src/confluence_check.jl
@@ -52,7 +52,7 @@ function check_confluence(
     is_reduced = isreduced(rws),
 )
     W = word_type(rws)
-    stack = Vector{Tuple{W,W}}()
+    stack = Stack{W}()
     return check_confluence!(stack, rws; is_reduced = is_reduced)
 end
 
diff --git a/src/knuthbendix2.jl b/src/knuthbendix2.jl
index 51abfc4..adf6ff8 100644
--- a/src/knuthbendix2.jl
+++ b/src/knuthbendix2.jl
@@ -170,7 +170,7 @@ function knuthbendix!(
     rws::RewritingSystem{W},
 ) where {W}
     work = Workspace(rws, settings)
-    stack = Vector{Tuple{W,W}}()
+    stack = Stack{W}()
     if !isreduced(rws)
         rws = reduce!(settings.algorithm, rws, work) # we begin with a reduced system
     end
diff --git a/src/knuthbendix_base.jl b/src/knuthbendix_base.jl
index 22c6493..ed1d3a7 100644
--- a/src/knuthbendix_base.jl
+++ b/src/knuthbendix_base.jl
@@ -96,4 +96,5 @@ function _kb_progress(prog::Progress, total, current, on_stack)
     return prog
 end
 
-word_type(stack::AbstractVector{<:Tuple{W,W}}) where {W<:AbstractWord} = W
+word_type(st::Stack) = word_type(typeof(st))
+word_type(::Type{<:Stack{W}}) where {W} = W
diff --git a/src/knuthbendix_delete.jl b/src/knuthbendix_delete.jl
index 4048c72..4962e0c 100644
--- a/src/knuthbendix_delete.jl
+++ b/src/knuthbendix_delete.jl
@@ -21,7 +21,7 @@ function knuthbendix!(
     rws::RewritingSystem{W},
 ) where {W}
     work = Workspace(rws, settings)
-    stack = Vector{Tuple{W,W}}()
+    stack = Stack{W}()
     rws = reduce!(settings.algorithm, rws, work) # we begin with a reduced system
 
     rwrules = __rawrules(rws)
diff --git a/src/knuthbendix_idxA.jl b/src/knuthbendix_idxA.jl
index 2dfa2e0..8c2cb55 100644
--- a/src/knuthbendix_idxA.jl
+++ b/src/knuthbendix_idxA.jl
@@ -68,7 +68,7 @@ function knuthbendix!(
     idxA::IndexAutomaton = IndexAutomaton(rws),
 ) where {W}
     @assert isreduced(rws)
-    stack = Vector{Tuple{W,W}}()
+    stack = Stack{W}()
     rwrules = __rawrules(rws)
     settings = work.settings
 
@@ -149,7 +149,8 @@ end
 
 function __post!(rws::AbstractRewritingSystem, rewriting, work::Workspace)
     settings = work.settings
-    stack = Vector{Tuple{word_type(rws),word_type(rws)}}()
+    W = word_type(rws)
+    stack = Stack{W}()
 
     if work.dropped_rules > 0
         @assert isempty(stack)
diff --git a/src/rewriting_system.jl b/src/rewriting_system.jl
index bc3fead..8f7176e 100644
--- a/src/rewriting_system.jl
+++ b/src/rewriting_system.jl
@@ -124,7 +124,7 @@ function RewritingSystem(
 end
 
 function RewritingSystem(
-    rwrules::Vector{Tuple{W,W}},
+    rwrules::Stack{W},
     order::RewritingOrdering;
     confluent::Bool = false,
     reduced::Bool = false,
diff --git a/src/rules.jl b/src/rules.jl
index af93aac..e5be9c9 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -51,6 +51,8 @@ Base.length(r::Rule) = 2
 Base.last(r::Rule) = first(iterate(r, 1))
 Base.eltype(::Type{Rule{W}}) where {W} = W
 
+word_type(::Type{<:Rule{W}}) where {W} = W
+
 Base.show(io::IO, r::Rule) = ((a, b) = r; print(io, a, " ⇒ ", b))
 
 """
diff --git a/src/settings_workspace.jl b/src/settings_workspace.jl
index 752f55b..69ad9a2 100644
--- a/src/settings_workspace.jl
+++ b/src/settings_workspace.jl
@@ -81,7 +81,7 @@ mutable struct Workspace{CA,T,H,S<:Settings{CA}}
     settings::S
     confluence_timer::Int
     dropped_rules::Int
-    dropped_stack::Vector{Tuple{Word{T},Word{T}}}
+    dropped_stack::Stack{Word{T}}
 end
 
 function Workspace(word_t, history, settings::Settings)
@@ -92,7 +92,7 @@ function Workspace(word_t, history, settings::Settings)
         settings,
         0,
         0,
-        Tuple{word_t,word_t}[],
+        Stack{word_t}(),
     )
 end