From 5c00b031ee32ae1820d094404a7277ac13a0b9b0 Mon Sep 17 00:00:00 2001
From: Venkateshprasad <32921645+ven-k@users.noreply.github.com>
Date: Sun, 4 Feb 2024 18:45:16 +0530
Subject: [PATCH] feat: add `WriteOnceReadMany` and utils

`UNIT_SYMBOLS`, `UNIT_VALUES`, `UNIT_MAPPING`, `ALL_VALUES`,`ALL_SYMBOLS`, ALL_MAPPING`, `SYMBOLIC_UNIT_VALUES` are instances of this collection type.
With this data type only a certain set of operations are permitted on these collections.
---
 src/DynamicQuantities.jl    |  5 ++++-
 src/register_units.jl       | 39 ++++++++++++++++++++--------------
 src/symbolic_dimensions.jl  | 18 ++++++++--------
 src/units.jl                |  9 ++++----
 src/write_once_read_many.jl | 42 +++++++++++++++++++++++++++++++++++++
 5 files changed, 82 insertions(+), 31 deletions(-)
 create mode 100644 src/write_once_read_many.jl

diff --git a/src/DynamicQuantities.jl b/src/DynamicQuantities.jl
index bf81cd6d..bb5973d7 100644
--- a/src/DynamicQuantities.jl
+++ b/src/DynamicQuantities.jl
@@ -11,8 +11,11 @@ export ustrip, dimension
 export ulength, umass, utime, ucurrent, utemperature, uluminosity, uamount
 export uparse, @u_str, sym_uparse, @us_str, uexpand, uconvert
 
+const INDEX_TYPE = UInt16
+
 include("internal_utils.jl")
 include("fixed_rational.jl")
+include("write_once_read_many.jl")
 include("types.jl")
 include("utils.jl")
 include("math.jl")
@@ -39,7 +42,7 @@ using .Units: UNIT_SYMBOLS
 let _units_import_expr = :(using .Units: m, g)
     append!(
         _units_import_expr.args[1].args,
-        map(s -> Expr(:(.), s), filter(s -> s ∉ (:m, :g), UNIT_SYMBOLS))
+        Expr(:(.), s) for s in UNIT_SYMBOLS if s ∉ (:m, :g)
     )
     eval(_units_import_expr)
 end
diff --git a/src/register_units.jl b/src/register_units.jl
index d0722b5f..905726ae 100644
--- a/src/register_units.jl
+++ b/src/register_units.jl
@@ -3,14 +3,17 @@ import .SymbolicUnits:
     SymbolicDimensionsSingleton, SYMBOLIC_UNIT_VALUES, update_symbolic_unit_values!
 
 # Update the unit collections
-function update_unit_mapping(name, value, unit_mapping::Dict{Symbol,Int} = UNIT_MAPPING)
-    unit_mapping[name] = length(unit_mapping) + 1
-end
+const UNIT_UPDATE_LOCK = Threads.SpinLock()
 
 function update_all_values(name_symbol, unit)
-    push!(ALL_SYMBOLS, name_symbol)
-    push!(ALL_VALUES, unit)
-    ALL_MAPPING[name_symbol] = INDEX_TYPE(length(ALL_MAPPING) + 1)
+    lock(UNIT_UPDATE_LOCK) do
+        push!(ALL_SYMBOLS, name_symbol)
+        push!(ALL_VALUES, unit)
+        i = lastindex(ALL_VALUES)
+        ALL_MAPPING[name_symbol] = i
+        UNIT_MAPPING[name_symbol] = i
+        update_symbolic_unit_values!(name_symbol)
+    end
 end
 
 # Register
@@ -20,16 +23,20 @@ end
 
 function _register_unit(name::Symbol, value)
     name_symbol = Meta.quot(name)
-    reg_expr = _lazy_register_unit(name, value)
-    push!(reg_expr.args,
-        quote
-            $update_unit_mapping($name_symbol, $value)
+    index = get(ALL_MAPPING, name, INDEX_TYPE(0))
+    if iszero(index)
+        reg_expr = _lazy_register_unit(name, value)
+        push!(reg_expr.args, quote
             $update_all_values($name_symbol, $value)
-            $update_symbolic_unit_values!($name_symbol)
-            # suppress the print of `SYMBOLIC_UNIT_VALUES`
             nothing
-        end
-    )
-    return reg_expr
+        end)
+        return reg_expr
+    else
+        unit = ALL_VALUES[index]
+        # When a utility function to expand `value` to its final form becomes
+        # available, enable the following check. This will avoid throwing an error
+        # if user is trying to register an existing unit with matching values.
+        # unit.value != value && throw("Unit $name is already defined as $unit")
+        throw("Unit `$name` is already defined as `$unit`")
+    end
 end
-
diff --git a/src/symbolic_dimensions.jl b/src/symbolic_dimensions.jl
index 4f9c184d..84f2069c 100644
--- a/src/symbolic_dimensions.jl
+++ b/src/symbolic_dimensions.jl
@@ -1,3 +1,5 @@
+import ..WriteOnceReadMany
+import ..INDEX_TYPE
 import .Units: UNIT_SYMBOLS, UNIT_MAPPING, UNIT_VALUES
 import .Constants: CONSTANT_SYMBOLS, CONSTANT_MAPPING, CONSTANT_VALUES
 
@@ -8,13 +10,9 @@ disambiguate_symbol(s) = s in SYMBOL_CONFLICTS ? Symbol(s, :_constant) : s
 # Prefer units over constants:
 # For example, this means we can't have a symbolic Planck's constant,
 # as it is just "hours" (h), which is more common.
-const INDEX_TYPE = UInt16
-# Prefer units over constants:
-# For example, this means we can't have a symbolic Planck's constant,
-# as it is just "hours" (h), which is more common.
-const ALL_SYMBOLS = [UNIT_SYMBOLS..., disambiguate_symbol.(CONSTANT_SYMBOLS)...]
-const ALL_VALUES = [UNIT_VALUES..., CONSTANT_VALUES...]
-const ALL_MAPPING = Dict(ALL_SYMBOLS .=> (INDEX_TYPE(1):INDEX_TYPE(length(ALL_SYMBOLS))))
+const ALL_SYMBOLS = WriteOnceReadMany([UNIT_SYMBOLS..., disambiguate_symbol.(CONSTANT_SYMBOLS)...])
+const ALL_VALUES = WriteOnceReadMany([UNIT_VALUES..., CONSTANT_VALUES...])
+const ALL_MAPPING = WriteOnceReadMany(Dict(s => INDEX_TYPE(i) for (i, s) in enumerate(ALL_SYMBOLS)))
 
 """
     AbstractSymbolicDimensions{R} <: AbstractDimensions{R}
@@ -375,6 +373,7 @@ module SymbolicUnits
     import ..DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE
     import ..DEFAULT_VALUE_TYPE
     import ..DEFAULT_DIM_BASE_TYPE
+    import ..WriteOnceReadMany
 
     # Lazily create unit symbols (since there are so many)
     module Constants
@@ -403,7 +402,7 @@ module SymbolicUnits
     import .Constants as SymbolicConstants
     import .Constants: SYMBOLIC_CONSTANT_VALUES
 
-    const SYMBOLIC_UNIT_VALUES = DEFAULT_SYMBOLIC_QUANTITY_TYPE[]
+    const SYMBOLIC_UNIT_VALUES = WriteOnceReadMany{Vector{DEFAULT_SYMBOLIC_QUANTITY_TYPE}}()
 
     function update_symbolic_unit_values!(unit, symbolic_unit_values = SYMBOLIC_UNIT_VALUES)
         @eval begin
@@ -415,7 +414,8 @@ module SymbolicUnits
         end
     end
 
-    update_symbolic_unit_values!.(UNIT_SYMBOLS)
+    update_symbolic_unit_values!(w::WriteOnceReadMany) = update_symbolic_unit_values!.(w._raw_data)
+    update_symbolic_unit_values!(UNIT_SYMBOLS)
 
     """
         sym_uparse(raw_string::AbstractString)
diff --git a/src/units.jl b/src/units.jl
index 0bc9743d..3cdba464 100644
--- a/src/units.jl
+++ b/src/units.jl
@@ -1,14 +1,14 @@
 module Units
 
+import ..WriteOnceReadMany
 import ..DEFAULT_DIM_TYPE
 import ..DEFAULT_VALUE_TYPE
 import ..DEFAULT_QUANTITY_TYPE
 
 @assert DEFAULT_VALUE_TYPE == Float64 "`units.jl` must be updated to support a different default value type."
 
-const UNIT_SYMBOLS = Symbol[]
-const UNIT_VALUES = DEFAULT_QUANTITY_TYPE[]
-const UNIT_MAPPING = Dict{Symbol,Int}()
+const UNIT_SYMBOLS = WriteOnceReadMany{Vector{Symbol}}()
+const UNIT_VALUES = WriteOnceReadMany{Vector{DEFAULT_QUANTITY_TYPE}}()
 
 macro _lazy_register_unit(name, value)
     return esc(_lazy_register_unit(name, value))
@@ -22,7 +22,6 @@ end
 function _lazy_register_unit(name::Symbol, value)
     name_symbol = Meta.quot(name)
     quote
-        haskey($UNIT_MAPPING, $name_symbol) && throw("Unit $($name_symbol) already exists.")
         const $name = $value
         push!($UNIT_SYMBOLS, $name_symbol)
         push!($UNIT_VALUES, $name)
@@ -207,6 +206,6 @@ end
 # The user should define these instead.
 
 # Update `UNIT_MAPPING` with all internally defined unit symbols.
-merge!(UNIT_MAPPING, Dict(UNIT_SYMBOLS .=> 1:lastindex(UNIT_SYMBOLS)))
+const UNIT_MAPPING = WriteOnceReadMany(Dict(s => i for (i, s) in enumerate(UNIT_SYMBOLS)))
 
 end
diff --git a/src/write_once_read_many.jl b/src/write_once_read_many.jl
new file mode 100644
index 00000000..5245b968
--- /dev/null
+++ b/src/write_once_read_many.jl
@@ -0,0 +1,42 @@
+"""
+    WriteOnceReadMany()
+
+Used for storing units, values, symbolic-units.
+"""
+struct WriteOnceReadMany{V}
+    _raw_data::V
+
+    WriteOnceReadMany(_raw_data) = new{typeof(_raw_data)}(_raw_data)
+    WriteOnceReadMany{T}() where T = WriteOnceReadMany(T())
+end
+
+# Utility functions
+for f in (:enumerate, :length, :lastindex)
+    @eval begin
+        Base.$f(w::WriteOnceReadMany) = $f(w._raw_data)
+    end
+end
+
+Base.getindex(w::WriteOnceReadMany, i::Union{Int, INDEX_TYPE,  Symbol}) = getindex(w._raw_data, i)
+
+Base.iterate(w::WriteOnceReadMany) = iterate(w._raw_data)
+Base.iterate(w::WriteOnceReadMany, i::Int) = iterate(w._raw_data, i)
+
+Base.intersect(w::WriteOnceReadMany, v::AbstractSet) = intersect(w._raw_data, v)
+Base.intersect(v::AbstractSet, w::WriteOnceReadMany) = intersect(v, w._raw_data)
+
+Base.push!(w::WriteOnceReadMany, val...) = push!(w._raw_data, val...)
+
+for f in (:findfirst, :filter)
+    @eval begin
+        Base.$f(val::Function, w::WriteOnceReadMany) = $f(val, w._raw_data)
+    end
+end
+
+Base.setindex!(w::DynamicQuantities.WriteOnceReadMany{Dict{Symbol, INDEX_TYPE}}, i::Int, s::Symbol) = setindex!(w, INDEX_TYPE(i), s)
+function Base.setindex!(w::DynamicQuantities.WriteOnceReadMany{Dict{Symbol, T}}, i::T, s::Symbol) where T <: Union{Int, INDEX_TYPE}
+    haskey(w._raw_data, s) && throw("Unit $s already exists at index $(w[s])")
+    setindex!(w._raw_data, i, s)
+end
+
+Base.get(w::WriteOnceReadMany{Dict{Symbol, INDEX_TYPE}}, a, b) = get(w._raw_data, a, b)