From 3a95493fb4fa3347a03ce42e9837cc752438f535 Mon Sep 17 00:00:00 2001
From: timholy <tim.holy@gmail.com>
Date: Fri, 8 Aug 2014 05:34:00 -0500
Subject: [PATCH] Better JLD support for arrays-of-arrays of basic types

---
 src/jld.jl   | 34 ++++++++++++++++++++++++++++++++++
 src/plain.jl |  1 +
 test/jld.jl  |  3 +++
 3 files changed, 38 insertions(+)

diff --git a/src/jld.jl b/src/jld.jl
index d459cf1e2..fa582d566 100644
--- a/src/jld.jl
+++ b/src/jld.jl
@@ -345,6 +345,16 @@ end
 read{T<:ByteString}(obj::JldDataset, ::Type{Array{T}}) = read(obj.plain, Array{T})
 read{T<:BitsKindOrByteString,N}(obj::JldDataset, ::Type{Array{T,N}}) = read(obj, Array{T})
 
+# Arrays-of-arrays of basic types
+function read{T<:HDF5BitsKind,M,N}(obj::JldDataset, ::Type{Array{Array{T,N},M}})
+    A = read(obj.plain, HDF5.HDF5Vlen{T})
+    if isempty(A) && exists(obj, "dims")
+        dims = a_read(obj.plain, "dims")
+        A = reshape(A, dims...)
+    end
+    convert(Array{Array{T,N},M}, A)
+end
+
 # Nothing
 read(obj::JldDataset, ::Type{Nothing}) = nothing
 read(obj::JldDataset, ::Type{Bool}) = bool(read(obj, Uint8))
@@ -515,6 +525,30 @@ end
 write{T<:Union(HDF5BitsKind, ByteString)}(parent::Union(JldFile, JldGroup), name::ByteString, data::Union(T, Array{T})) =
     write(parent, name, data, full_typename(typeof(data)))
 
+# Arrays-of-arrays of basic types
+write{T<:Union(HDF5BitsKind, ByteString),N}(parent::Union(JldFile, JldGroup), name::ByteString,
+                                            data::Array{Array{T,N}}, astype::ByteString) = 
+    write(parent, name, HDF5.HDF5Vlen(data), astype)
+write{T<:Union(HDF5BitsKind, ByteString),N}(parent::Union(JldFile, JldGroup), name::ByteString,
+                                            data::Array{Array{T,N}}) =
+    write(parent, name, data, full_typename(typeof(data)))
+function write{T}(parent::Union(JldFile, JldGroup), name::ByteString,
+                  data::HDF5.HDF5Vlen{T}, astype::ByteString)
+    # Create the dataset
+    dset, dtype = d_create(parent.plain, name, data)
+    try
+        # Write the attribute
+        a_write(dset, name_type_attr, astype)
+        isa(data, Array) && isempty(data) && a_write(dset, "dims", [size(data)...])
+        # Write the data
+        HDF5.writearray(dset, dtype.id, data)
+    finally
+        close(dset)
+        close(dtype)
+    end
+end
+
+
 # Write nothing
 function write(parent::Union(JldFile, JldGroup), name::ByteString, n::Nothing, astype::ASCIIString)
     local dspace, dset
diff --git a/src/plain.jl b/src/plain.jl
index 9ca0d7a23..c6275799b 100644
--- a/src/plain.jl
+++ b/src/plain.jl
@@ -419,6 +419,7 @@ type HDF5Vlen{T}
 end
 HDF5Vlen{S<:ByteString}(strs::Array{S}) = HDF5Vlen{chartype(S)}(strs)
 HDF5Vlen{T<:HDF5BitsKind}(A::Array{Array{T}}) = HDF5Vlen{T}(A)
+HDF5Vlen{T<:HDF5BitsKind,N}(A::Array{Array{T,N}}) = HDF5Vlen{T}(A)
 
 ## Types that correspond to C structs and get used for ccall
 # For VLEN
diff --git a/test/jld.jl b/test/jld.jl
index 38db809e7..7b22b0bb8 100644
--- a/test/jld.jl
+++ b/test/jld.jl
@@ -5,6 +5,7 @@ using Base.Test
 # Define variables of different types
 x = 3.7
 A = rand(1:20, 3, 5)
+Aarray = Vector{Float64}[[1.2,1.3],[2.2,2.3,2.4]]
 str = "Hello"
 stringsA = ASCIIString["It", "was", "a", "dark", "and", "stormy", "night"]
 stringsU = UTF8String["It", "was", "a", "dark", "and", "stormy", "night"]
@@ -141,6 +142,7 @@ fn = joinpath(tempdir(),"test.jld")
 fid = jldopen(fn, "w")
 @write fid x
 @write fid A
+@write fid Aarray
 @write fid str
 @write fid stringsA
 @write fid stringsU
@@ -196,6 +198,7 @@ for mmap = (true, false)
     fidr = jldopen(fn, "r", mmaparrays=mmap)
     @check fidr x
     @check fidr A
+    @check fidr Aarray
     @check fidr str
     @check fidr stringsA
     @check fidr stringsU