From 99cd250b630def8820ebfe09b881643af5c96908 Mon Sep 17 00:00:00 2001
From: Marek Kaluba <kalmar@amu.edu.pl>
Date: Sat, 18 Dec 2021 16:40:41 +0100
Subject: [PATCH] add WeightedLex

---
 src/orderings.jl  | 53 ++++++++++++++++++++++++++++++++++++++++++++++-
 test/orderings.jl | 11 ++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/src/orderings.jl b/src/orderings.jl
index 585d481d..075e09b8 100644
--- a/src/orderings.jl
+++ b/src/orderings.jl
@@ -1,5 +1,5 @@
 import Base.Order: lt, Ordering
-export LenLex, WreathOrder, RecursivePathOrder
+export LenLex, WreathOrder, RecursivePathOrder, WeightedLex
 
 """
     WordOrdering <: Ordering
@@ -122,3 +122,54 @@ function lt(o::RecursivePathOrder, p::AbstractWord, q::AbstractWord)
     end
     return false
 end
+
+"""
+    struct WeightedLex{T} <: WordOrdering
+    WeightedLex(A::Alphabet, weights::AbstractVector)
+
+`WeightedLex` order compares words first according to their weight and then
+by the lexicographic order determined by the order of letters `A`.
+The `weight` array assigns weights to each letter and the weight of a word is
+simply the sum of weights of all letters.
+The `LenLex` ordering is a special case of `WeightedLex` when all weights are equal to `1`.
+
+!!! note:
+    Since empty word is assigned a value of `zero(eltype(weights))` a vector of
+    positive weights is strongly recommended.
+"""
+struct WeightedLex{T,S} <: WordOrdering
+    A::Alphabet{T}
+    weights::Vector{S}
+
+    function WeightedLex(A::Alphabet{T}, weights::AbstractVector{S}) where {T,S}
+        @assert length(weights) == length(A)
+        @assert all(w-> w >=(zero(S)), weights)
+        return new{T,S}(A, weights)
+    end
+end
+
+Base.hash(wl::WeightedLex, h::UInt) = hash(wl.weights, hash(wl.lenlex, h))
+
+Base.@propagate_inbounds weight(o::WeightedLex, l::Integer) = o.weights[l]
+
+function weight(o::WeightedLex, p::AbstractWord)
+    isone(p) && return zero(eltype(o.weights))
+    return @inbounds sum(weight(o, l) for l in p)
+end
+
+function Base.Order.lt(o::WeightedLex, p::AbstractWord, q::AbstractWord)
+    S = eltype(o.weights)
+
+    weight_p = weight(o, p)
+    weight_q = weight(o, q)
+
+    if weight_p == weight_q
+        for (a, b) in zip(p, q)
+            # comparing only on positive pointer values
+            a == b || return isless(a, b)
+        end
+        return false
+    else
+        return isless(weight_p, weight_q)
+    end
+end
diff --git a/test/orderings.jl b/test/orderings.jl
index 400fdaa8..17da81dc 100644
--- a/test/orderings.jl
+++ b/test/orderings.jl
@@ -55,4 +55,15 @@
     @test sort(a, order = rpo) == [ε, w1, w14, w214, w1224]
     @test sort(b, order = rpo) == [ε, w1, w2, w214, w41, w241, w32141]
 
+    @test_throws AssertionError WeightedLex(A, [1,2,3])
+    @test_throws AssertionError WeightedLex(A, [1,2, 3,-4])
+
+    wtlex = WeightedLex(A, [1,2,3,4])
+
+    @test lt(wtlex, one(Word{Int}), Word([1]))
+    @test !lt(wtlex, Word([1]), Word([1]))
+    @test lt(wtlex, Word([1,1]), Word([2]))
+    @test lt(wtlex, Word([2]), Word([1,1,1]))
+    @test lt(wtlex, Word([2,4]), Word([4,2]))
+    @test lt(wtlex, Word([2,4]), Word([4,1,1]))
 end