Julep: Efficient Hierarchical Mutable Data

I completed the (POC of the) macro-based version this morning. Now it works for nested structs:

abstract type AbstractElephant end

struct Elephant <: AbstractElephant
    hunger::Float64
end

@with_inline_types mutable struct Forest
    elephant::@inline(Elephant)   # @inline elephant::Elephant also possible, but less
                                  # MacroTools-friendly
    normal_field::String
end

@with_inline_types mutable struct Biome
    forest::@inline(Forest)
end

julia> b = Biome(Forest(Elephant(50), "Tasmania"))
Biome(50.0, "Tasmania")

julia> b.forest.elephant.hunger = 33.0
33.0

julia> dump(b)
Biome
  forest__elephant__hunger: Float64 33.0
  forest__normal_field: String "Tasmania"

julia> @btime $b.forest.elephant.hunger
  2.768 ns (0 allocations: 0 bytes)
33.0

I used generated functions to limit the complexity of the macro expansion. The way it’s written, it might also not be that difficult to support parametric types…

PatrickHaecker Would that meet your needs?

I get the appeal of getting a new keyword in base to solve such a problem, but as you probably are well aware, that does not happen very often. Pushing this macro to the finishing line is more likely to work out… Happy to help if there is interest.

Full code:

using MacroTools
using MacroTools: @qq
using Base.Meta: quot

""" Defined for types, not objects. """
static_propertynames(::Type{T}) where T = fieldnames(T)

get_inlined_type_fields(x) = Val((;))

@generated function get_view_property(::Val{inl_typ_fields}, view::T, ::Val{p2}) where {T, p2, inl_typ_fields}
    # When evaluating `b.forest.elephant.hunger`, this is called twice. First time O is
    # ForestView{:forest, Biome}, and we return
    #     ElephantView{:forest__elephant, Biome}(view.obj)
    # (simplified) from the first branch below while second time we return
    #     view.obj.forest__elephant__hunger
    # from the second branch. The idea being that each nested View object has, as
    # parameter, the corresponding field prefix.
    @assert length(T.parameters) == 2
    p1 = T.parameters[1]
    obj_expr = Expr(:call, :getfield, :view, quot(:obj))
    if haskey(inl_typ_fields, p2)
        sub_view = inl_typ_fields[p2]
        p12 = quot(Symbol(p1, "__", p2))
        return :($sub_view{$p12, typeof($obj_expr)}($obj_expr))
    else
        return Expr(:call, :getfield, obj_expr, quot(Symbol(p1, "__", p2)))
    end
end

@generated function set_view_property!(::Val{inl_typ_fields}, view::T, ::Val{p2}, new_val) where {T, p2, inl_typ_fields}
    @assert length(T.parameters) == 2
    p1 = T.parameters[1]
    obj_expr = Expr(:call, :getfield, :view, quot(:obj))
    if haskey(inl_typ_fields, p2)
        # Eg. this happens for
        #     b.forest.elephant = Elephant(...)
        error("TODO: must copy all fields")
    else
        # TODO: put a `convert` call to the right type
        return Expr(:call, :setfield!, obj_expr, quot(Symbol(p1, "__", p2)), :new_val)
    end
end

macro with_inline_types(def)
    # Example with respect to this
    #     @with_inline_types struct Forest
    #         elephant::@inline(Elephant)
    #         normal_field::String
    #     end

    di = MacroTools.splitstructdef(def)
    di2 = copy(di)
    name = di[:name]  # Forest
    inline_code = Expr[]  # code to define the type views
    di2[:fields] = []     # [(normal_field, Int)]
    inlined_type_fields = Dict()  # [:elephant=>ElephantView]
    prop_names = Symbol[]   # [:elephant, :normal_field]
    constr_accessors = []

    for (f, type) in di[:fields]
        push!(prop_names, f)
        if @capture(type, @inline(Texpr_))
            T = eval(Texpr)  # having to `eval` is unfortunate. there are other ways worth
                             # investigating, such as `getfield(@__MODULE__, Texpr)`, that would
                             # have different trade-offs
            s_fields = [Symbol(f, "__", sub_f) for sub_f in fieldnames(T)] # [:elephant__hunger]
            view_name = Symbol(Texpr, "View")
            type_view_code = quote
                struct $view_name{F, U} <: $(supertype(T)) # ElephantView <: AbstractElephant
                    obj::U
                end
                Base.propertynames(view::$view_name) = $(fieldnames(T))
                Base.getproperty(view::$view_name, p::Symbol) =
                    get_view_property(get_inlined_type_fields($Texpr), view, Val{p}())
                Base.setproperty!(view::$view_name, p::Symbol, x) =
                    set_view_property!(get_inlined_type_fields($Texpr), view, Val{p}(), x)
            end
            for s_field in fieldnames(T)
                push!(constr_accessors, :(getfield($f, $(quot(s_field)))))
            end
            push!(inline_code, type_view_code)
            inlined_type_fields[f] = :($view_name)
            append!(di2[:fields], map(tuple, s_fields, fieldtypes(T)))
        else
            push!(constr_accessors, f)
            push!(di2[:fields], (f, type))
        end
    end
    @qq(begin  # @qq to get good line numbers
        $(MacroTools.combinestructdef(di2))
        $(inline_code...)
        function Base.getproperty(obj::$name, p::Symbol)
            # if p == :elephant
            #      return ElephantViewFromForest(obj)
            # end
            # return getfield(obj, p)
            $([:(if p === $(quot(fieldname))
                 return $view_type{$(quot(fieldname)), $name}(obj)
                 end)
               for (fieldname, view_type) in inlined_type_fields]...)
            return Base.getfield(obj, p)
        end
        Base.propertynames(obj::$name) = $prop_names
        static_propertynames(::Type{$name}) = $prop_names
        function $name($(prop_names...))
             $name($(constr_accessors...))
        end
        get_inlined_type_fields(::Type{$name}) = $(Val(NamedTuple(inlined_type_fields)))
        $name
    end) |> esc
end
3 Likes