feat: Materialize ucs4/ascii array data as string views

Author: icweaver

src/ASDF.jl

@@ -492,10 +492,12 @@ An ASDF fixed-length ASCII string datatype, corresponding to the `["ascii", N]`
 the ASDF datatype spec. Each element occupies exactly `length` bytes, one byte per
 character (all codepoints < 128).
 
-The corresponding Julia type is `NTuple{N, UInt8}`. Byte order is irrelevant for this
-type since each character is a single byte.
+The corresponding `isbitstype` Julia type is `NTuple{N, UInt8}`. Byte order is irrelevant for
+this type since each character is a single byte. When an array of this datatype is materialized
+with [`Base.getindex(ndarray::NDArray)`](@ref), each element is presented as an
+[`AsciiString`](@ref): a thin `AbstractString` view over those bytes, with identical contents.
 
-See also: [`Ucs4Datatype`](@ref), [`parse_asdf_datatype`](@ref).
+See also: [`Ucs4Datatype`](@ref), [`AsciiString`](@ref), [`parse_asdf_datatype`](@ref).
 """
 struct AsciiDatatype
     length::Int  # Bytes (1 per char)
@@ -508,9 +510,12 @@ An ASDF fixed-length UCS-4 string datatype, corresponding to the `["ucs4", N]` f
 the ASDF datatype spec. Each element occupies exactly `4 * length` bytes, with each
 character encoded as a 4-byte UInt32 in the array's declared byte order.
 
-The corresponding Julia type is `NTuple{N, UInt32}`.
+The corresponding `isbitstype` Julia type is `NTuple{N, UInt32}`. When an array of this
+datatype is materialized with [`Base.getindex(ndarray::NDArray)`](@ref), each element is
+presented as a [`UCS4String`](@ref): a thin `AbstractString` view over those codepoints that
+displays and behaves as a Julia string, with identical bytes.
 
-See also: [`AsciiDatatype`](@ref), [`parse_asdf_datatype`](@ref).
+See also: [`AsciiDatatype`](@ref), [`UCS4String`](@ref), [`parse_asdf_datatype`](@ref).
 """
 struct Ucs4Datatype
     length::Int  # Characters (4 bytes each)
@@ -788,9 +793,10 @@ Returns the fully materialized array. See [`ASDF.NDArray`](@ref) for definitions
 
 ```julia
 size(result) == Tuple(reverse(ndarray.shape))
-eltype(result) == Type(ndarray.datatype)
-sizeof(eltype) .* strides(result) == Tuple(reverse(ndarray.strides))
+eltype(result) == ASDF.materialized_eltype(ndarray.datatype)
 ```
+
+For the `ucs4` and `ascii` string datatypes, [`materialized_eltype`](@ref) is a thin `AbstractString` view over the characters ([`UCS4String`](@ref) / [`AsciiString`](@ref); see [`stringify_data`](@ref)). For all other datatypes, `eltype(result) == Type(ndarray.datatype)` and additionally `sizeof(eltype) .* strides(result) == Tuple(reverse(ndarray.strides))`.
 """
 function Base.getindex(ndarray::NDArray)
     if ndarray.data !== nothing
@@ -831,6 +837,12 @@ function Base.getindex(ndarray::NDArray)
         error("`data` has different stride from `ndarray.strides`")
     end
 
+    # Present `ucs4`/`ascii` data as `UCS4String`/`AsciiString` rather than raw tuples of
+    # codepoints. Done after the layout checks above, which rely on the `isbitstype` tuple
+    # representation; the views are a zero-copy `reinterpret` sharing those bytes, so this is
+    # display/behavior only.
+    data = stringify_data(data, ndarray.datatype)
+
     return data::AbstractArray
 end
 
@@ -883,6 +895,113 @@ function correct_byteorder(data, ::Ucs4Datatype, byteorder::Byteorder)
     return data
 end
 
+"""
+    UCS4String{N} <: AbstractString
+
+A fixed-width UCS-4 string of `N` characters, as materialized from an [`Ucs4Datatype`](@ref)
+array. It wraps the raw `NTuple{N, UInt32}` codepoints and is `isbitstype` with an identical
+memory layout, so it reinterprets to and from the on-disk block bytes exactly like the tuple
+would. The difference is purely in presentation: a `UCS4String` *displays* and *behaves* as a
+Julia string rather than as raw `UInt32` codepoints.
+
+```jldoctest
+julia> s = ASDF.UCS4String((UInt32('h'), UInt32('i'), UInt32(0)))  # null-padded
+"hi"
+
+julia> s == "hi"
+true
+```
+
+As a fixed-width string, trailing null (`\\0`) padding is treated as unused and excluded from
+the string's length and contents. The full `N` codepoints are retained internally so the
+on-disk representation round-trips unchanged.
+
+See also: [`AsciiString`](@ref).
+"""
+struct UCS4String{N} <: AbstractString
+    codes::NTuple{N, UInt32}
+end
+
+"""
+    AsciiString{N} <: AbstractString
+
+A fixed-width ASCII string of `N` characters, as materialized from an [`AsciiDatatype`](@ref)
+array. The single-byte analogue of [`UCS4String`](@ref): it wraps the raw `NTuple{N, UInt8}`
+bytes and is `isbitstype` with an identical memory layout, reinterpreting to and from the
+on-disk block bytes exactly like the tuple would, while *displaying* and *behaving* as a Julia
+string rather than as raw `UInt8` bytes.
+
+```jldoctest
+julia> s = ASDF.AsciiString((UInt8('h'), UInt8('i'), UInt8(0)))  # null-padded
+"hi"
+
+julia> s == "hi"
+true
+```
+
+As with `UCS4String`, trailing null (`\\0`) padding is excluded from the string's length and
+contents, while the full `N` bytes are retained internally so the on-disk representation
+round-trips unchanged.
+"""
+struct AsciiString{N} <: AbstractString
+    codes::NTuple{N, UInt8}
+end
+
+# Both ASDF string views share the same fixed-width, null-padded, one-codeunit-per-character
+# `AbstractString` interface; only the code-unit type differs.
+const ASDFString = Union{UCS4String, AsciiString}
+
+# Number of significant characters, i.e. `N` minus any trailing null padding.
+function significant_length(s::ASDFString)
+    n = length(s.codes)
+    @inbounds while n > 0 && s.codes[n] == 0
+        n -= 1
+    end
+    return n
+end
+
+Base.ncodeunits(s::ASDFString) = significant_length(s)
+Base.codeunit(::UCS4String) = UInt32
+Base.codeunit(::AsciiString) = UInt8
+function Base.codeunit(s::ASDFString, i::Integer)
+    @boundscheck 1 <= i <= ncodeunits(s) || throw(BoundsError(s, i))
+    return @inbounds s.codes[i]
+end
+Base.isvalid(s::ASDFString, i::Integer) = 1 <= i <= ncodeunits(s)
+Base.length(s::ASDFString) = ncodeunits(s)
+function Base.iterate(s::ASDFString, i::Int = 1)
+    # State is the code-unit index, per the `AbstractString` interface (`getindex` calls
+    # `iterate(s, i::Integer)`), so the significant length is rechecked each step.
+    i > ncodeunits(s) && return nothing
+    u = s.codes[i]
+    isvalid(Char, u) || error("$(typeof(s)) contains an invalid Unicode codepoint $(repr(u)) at index $i")
+    return (Char(u), i + 1)
+end
+
+"""
+    stringify_data(data, datatype) -> AbstractArray
+
+Present a materialized [`Ucs4Datatype`](@ref) array as [`UCS4String`](@ref)s, and an
+[`AsciiDatatype`](@ref) array as [`AsciiString`](@ref)s, instead of raw tuples of codepoints.
+For any other datatype the data is returned unchanged. The string views are `isbitstype` with
+the same layout as the backing tuples, so this is a zero-copy `reinterpret`; only display and
+behavior change, and writing is unaffected.
+"""
+stringify_data(data, datatype) = data
+stringify_data(data, dt::Ucs4Datatype) = reinterpret(materialized_eltype(dt), data)
+stringify_data(data, dt::AsciiDatatype) = reinterpret(materialized_eltype(dt), data)
+
+"""
+    materialized_eltype(datatype) -> Type
+
+Julia element type produced by materializing an array of the given ASDF `datatype`. This is
+[`UCS4String{N}`](@ref UCS4String) for [`Ucs4Datatype`](@ref), [`AsciiString{N}`](@ref AsciiString)
+for [`AsciiDatatype`](@ref), and `Type(datatype)` otherwise.
+"""
+materialized_eltype(dt) = Type(dt)
+materialized_eltype(dt::Ucs4Datatype) = UCS4String{dt.length}
+materialized_eltype(dt::AsciiDatatype) = AsciiString{dt.length}
+
 """
     TaggedMapping{D} <: AbstractDict
     TaggedSequence{V} <: AbstractVector
@@ -897,8 +1016,10 @@ struct TaggedMapping{D<:AbstractDict} <: AbstractDict{Any, Any}
     tag::String
     value::D
 end
-# `keys`, `values`, `haskey`, `get`, etc. fall back to these via the `AbstractDict` interface.
+# `keys`, `values`, `haskey`, etc. fall back to these via the `AbstractDict` interface.
+# `get` is delegated explicitly because the `haskey`/`get` fallbacks need a 3-arg `get`.
 Base.getindex(m::TaggedMapping, k) = getindex(m.value, k)
+Base.get(m::TaggedMapping, k, default) = get(m.value, k, default)
 Base.length(m::TaggedMapping) = length(m.value)
 Base.iterate(m::TaggedMapping, state...) = iterate(m.value, state...)
 
@@ -1060,7 +1181,7 @@ Allocates a dense array of shape `reverse(shape)` and fills it by calling `chunk
 """
 function Base.getindex(chunked_ndarray::ChunkedNDArray)
     shape = chunked_ndarray.shape
-    datatype = Type(chunked_ndarray.datatype)
+    datatype = materialized_eltype(chunked_ndarray.datatype)
     data = Array{datatype}(undef, reverse(shape)...)
     for chunk in chunked_ndarray.chunks
         start = CartesianIndex(reverse(chunk.start .+ 1)...)
@@ -1395,8 +1516,8 @@ const blocks::Blocks = Blocks()
 
 Infer the ASDF datatype from a Julia element type, used when writing arrays:
 
-- `NTuple{N, UInt8}` --> [`AsciiDatatype(N)`](@ref AsciiDatatype)
-- `NTuple{N, UInt32}` --> [`Ucs4Datatype(N)`](@ref Ucs4Datatype)
+- `NTuple{N, UInt8}` or [`AsciiString{N}`](@ref AsciiString) --> [`AsciiDatatype(N)`](@ref AsciiDatatype)
+- `NTuple{N, UInt32}` or [`UCS4String{N}`](@ref UCS4String) --> [`Ucs4Datatype(N)`](@ref Ucs4Datatype)
 - `NamedTuple` --> [`StructuredDatatype`](@ref) with fields inferred recursively
 - Any other type --> [`Datatype`](@ref) via the existing `type_datatype_dict` lookup
 
@@ -1419,6 +1540,11 @@ function infer_asdf_datatype(T::Type)::Union{Datatype, AsciiDatatype, Ucs4Dataty
     end
 end
 
+# `UCS4String{N}`/`AsciiString{N}` round-trip to the datatypes they were materialized from; the
+# width `N` is the type parameter, recovered directly by dispatch (the inverse of `materialized_eltype`).
+infer_asdf_datatype(::Type{UCS4String{N}}) where {N} = Ucs4Datatype(N)
+infer_asdf_datatype(::Type{AsciiString{N}}) where {N} = AsciiDatatype(N)
+
 function YAML._print(io::IO, val::NDArrayWrapper, level::Int=0, ignore_level::Bool=false)
     datatype = infer_asdf_datatype(eltype(val.array))
     if val.inline

test/test-datatypes.jl

@@ -0,0 +1,100 @@
+# Unit tests for the ASDF string/structured datatype helpers and the write-side
+# `infer_asdf_datatype` inference. The reference roundtrips exercise these indirectly,
+# but the methods below are pinned down here so they stay covered without large fixtures.
+
+@testset "string datatype interface" begin
+    ucs = ASDF.UCS4String((UInt32('h'), UInt32('i'), UInt32(0)))  # null-padded
+    asc = ASDF.AsciiString((UInt8('h'), UInt8('i'), UInt8(0)))
+
+    # `codeunit(::Type)` reports the code-unit type for each string view.
+    @test codeunit(ucs) == UInt32
+    @test codeunit(asc) == UInt8
+
+    # Indexed `codeunit` returns the raw code unit at that position.
+    @test codeunit(ucs, 1) == UInt32('h')
+    @test codeunit(ucs, 2) == UInt32('i')
+    @test codeunit(asc, 1) == UInt8('h')
+    @test codeunit(asc, 2) == UInt8('i')
+
+    # Trailing null padding is excluded from the significant length.
+    @test ncodeunits(ucs) == 2
+    @test ncodeunits(asc) == 2
+
+    # `isvalid` is bounded by the significant length, not the full tuple width.
+    @test isvalid(ucs, 1)
+    @test isvalid(ucs, 2)
+    @test !isvalid(ucs, 3)   # padding byte
+    @test !isvalid(ucs, 0)
+    @test isvalid(asc, 2)
+    @test !isvalid(asc, 3)
+
+    @test ucs == "hi"
+    @test asc == "hi"
+
+    # Embedded (non-trailing) nulls are kept: `significant_length` strips only *trailing*
+    # nulls, so the significant length runs through to the last non-null code unit.
+    ucs_embedded = ASDF.UCS4String((UInt32('h'), UInt32(0), UInt32('i')))
+    asc_embedded = ASDF.AsciiString((UInt8('h'), UInt8(0), UInt8('i')))
+    @test ncodeunits(ucs_embedded) == 3
+    @test ncodeunits(asc_embedded) == 3
+    @test ucs_embedded == "h\0i"
+    @test asc_embedded == "h\0i"
+end
+
+@testset "materialized_eltype" begin
+    @test ASDF.materialized_eltype(ASDF.Ucs4Datatype(3)) == ASDF.UCS4String{3}
+    @test ASDF.materialized_eltype(ASDF.AsciiDatatype(4)) == ASDF.AsciiString{4}
+    # Plain scalar datatypes fall through to the Julia type.
+    @test ASDF.materialized_eltype(ASDF.Datatype_float32) == Float32
+end
+
+@testset "infer_asdf_datatype" begin
+    # Raw `NTuple` byte/codepoint forms map to the ASDF string datatypes.
+    @test ASDF.infer_asdf_datatype(NTuple{2, UInt8}) == ASDF.AsciiDatatype(2)
+    @test ASDF.infer_asdf_datatype(NTuple{2, UInt32}) == ASDF.Ucs4Datatype(2)
+
+    # The string *views* recover the same datatypes as their backing tuples.
+    @test ASDF.infer_asdf_datatype(ASDF.AsciiString{5}) == ASDF.AsciiDatatype(5)
+    @test ASDF.infer_asdf_datatype(ASDF.UCS4String{6}) == ASDF.Ucs4Datatype(6)
+
+    # `NamedTuple` becomes a structured datatype with fields inferred recursively.
+    sd = ASDF.infer_asdf_datatype(@NamedTuple{x::Float32, label::NTuple{4, UInt8}})
+    @test sd isa ASDF.StructuredDatatype
+    @test [f.name for f in sd.fields] == ["x", "label"]
+    @test sd.fields[1].datatype == ASDF.Datatype_float32
+    @test sd.fields[2].datatype == ASDF.AsciiDatatype(4)
+
+    # Plain scalar types defer to the existing dict lookup.
+    @test ASDF.infer_asdf_datatype(Int64) == ASDF.Datatype_int64
+end
+
+@testset "tagged node accessors" begin
+    # The `Tagged*` wrappers delegate their collection / string interfaces to the wrapped
+    # value. Loading round-trips them (see test-read_fallbacks), but the individual delegated
+    # methods are pinned down directly here so they stay covered across platforms.
+    tm  = ASDF.TaggedMapping("tag:example.org:mylib/widget-1.0.0", Dict("a" => 1, "b" => 2))
+    ts  = ASDF.TaggedSequence("tag:example.org:mylib/series-1.0.0", ["alpha", "beta", "gamma"])
+    tsc = ASDF.TaggedScalar("tag:example.org:mylib/quantity-1.0.0", "3.14")
+
+    # TaggedMapping behaves as its underlying dict.
+    @test length(tm) == 2
+    @test tm["a"] == 1
+    @test haskey(tm, "b")
+    @test !haskey(tm, "missing")
+    @test get(tm, "b", 0) == 2
+    @test get(tm, "missing", -1) == -1
+
+    # TaggedSequence behaves as its underlying vector, with linear indexing.
+    @test Base.IndexStyle(typeof(ts)) == IndexLinear()
+    @test Base.IndexStyle(ASDF.TaggedSequence) == IndexLinear()
+    @test size(ts) == (3,)
+    @test ts[2] == "beta"
+
+    # TaggedScalar behaves as its underlying string.
+    @test ncodeunits(tsc) == ncodeunits("3.14")
+    @test codeunit(tsc) == UInt8
+    @test codeunit(tsc, 1) == codeunit("3.14", 1)
+    @test isvalid(tsc, 1)
+    @test !isvalid(tsc, 99)
+    @test tsc == "3.14"
+end

test/test-ndarray.jl

@@ -103,7 +103,9 @@ end
         @test nd[] == expected
 
         nd = make_ndarray(; lazy_block_headers = lbh,  source = Int64(0), data = nothing, byteorder = opposite, datatype = ASDF.Ucs4Datatype(2), shape = [Int64(1)], strides = [Int64(8)])
-        @test nd[] == [(UInt32(0x00000001), UInt32(0x00000002))]
+        # ucs4 data is presented as a string-like `UCS4String`, byte-swapped to host order
+        @test eltype(nd[]) <: AbstractString
+        @test nd[] == [join((Char(0x1), Char(0x2)))]
     end
 
     nd = make_ndarray(; strides = Int64[5])