Module with state - how to implement?

I would like to have a module with state: when I’m working in the repl, I load the module and its state is “empty”; as I use the module’s constructors to instantiate variables of a certain type defined in that module, the module should keep track of those variables (for example, by keeping a list of all such variables; you can assume variables of this type never go out of scope). My question is: how to implement this?

I know that each module has its global scope, so at first I thought that maybe the module could instantiate the container to keep track of variables instantiate by the module’s constructors, but I read in performance tips that you should avoid global variables.

Could you give an example of the module and what you’d like to track?

I think most of the performance can be gained by declaring the global as const

1 Like

That’s good advice. I would take this a step further: Don’t use a module with state, period.

Instead, create a struct or mutable struct to represent the state and store your variables in that struct. This has a wide variety of advantages:

  • It completely avoids the problem with performance of global variables
  • It allows you to maintain multiple states at the same time by just constructing multiple instances of the struct.
  • It completely avoids thread-safety issues where multiple parallel tasks try to modify the global state of the module at the same time.
  • It is more idiomatic and should be easier for other Julia programmers to understand when they read your code.
  • It allows variables to actually go out of scope when you are done with them. I know you’ve said we can assume they don’t, but it is actually nice to be able to reclaim memory eventually…
4 Likes

Well I don’t have it yet, I’m trying to think ahead before starting to code. But schematically:

module ModWithState

global_state = Dict()

struct TrackedVar{T}
    value::T
end

end

Now what I want to do and don’t know how, is to keep track of a kind of graph structure between these tracked vars. For example, if I were to type in the repl (schematically):

x = TrackedVar(5)
y = TrackedVar(2)
z = x + y

Then z should be a TrackedVar with value 7 and furthermore, global_state should be somehow encode the graph structure implied by z=x+y, for example something like Dict(x=>[], y=>[], z=>[x,y]).

I am aware of the things you’re pointing out. I’d like to hear your thoughts on the example I presented.

I wanted to reply to this, but my post was not counted as a reply, and I don’t know how to edit to make it so. But please see above.

I would suggest removing global_state entirely, and instead passing a “workspace” or “state” or “model”, to the TrackedVar constructor. Something like:

struct Model
  dependencies::Dict  # TODO: make this concretely typed
end

struct TrackedVar{T}
  value::T
end

function add_variable!(model::Model, value)
  var = TrackedVar(value)
  model.dependencies[var] = []
  var
end

You might also consider adding a name::Symbol to TrackedVar and providing that at construction. You could even create a macro like:

@variable model x = 5

which would produce something like:

x = add_variable!(model, 5, :x)

Note that this is quite similar to the way variables work in JuMP.jl, which is generally a good model to emulate.

1 Like

This is roughly what I currently have (minus the macro) :slight_smile:

But I still think there would be value in being able to write z=x+y without add_variable!. I do like the macro idea though. If I can’t have my way and use module state then I’ll implement that :wink:

You should generally avoid carrying around global state and do something like @rdeits example instead, but if you absolutely have to do it like that, you can achieve this with something like:

module ModWithState

const global_state = IdDict()

mutable struct TrackedVar{T}
    value::T
    function TrackedVar(x::T, deps=[]) where {T}
        t = new{T}(x)
        global_state[t] = deps
        return t
    end
end

function Base.:+(x::TrackedVar, y::TrackedVar)
    return TrackedVar(x.value + y.value, [x, y])
end

end
julia> using .ModWithState: TrackedVar

julia> x = TrackedVar(5)
TrackedVar{Int64}(5)

julia> y = TrackedVar(2)
TrackedVar{Int64}(2)

julia> z = x + y
TrackedVar{Int64}(7)

julia> ModWithState.global_state
IdDict{Any,Any} with 3 entries:
  TrackedVar{Int64}(2) => Any[]
  TrackedVar{Int64}(7) => TrackedVar{Int64}[TrackedVar{Int64}(5), TrackedVar{Int64}(2)]
  TrackedVar{Int64}(5) => Any[]

Two things to keep in mind though:

  1. With this implementation, TrackedVars will never get garbage collected, so if you are creating a lot of them, you might run out of memory. The way to work around this would probably be to only store WeakRefs in the global dict.
  2. Creating TrackedVars this way will not be thread safe, so you can’t create TrackedVars from different threads without creating potential race conditions. To fix this, you would probably need to carry around a global lock for global_state.

Edit: You probably want to use an IdDict here, instead of a Dict.

2 Likes

Thanks for the answer! I don’t really have to do it like that, I just believe it’s nicer for the user.

Now that I have you here, here’s a follow up. You defined Base.:+(x::TrackedVar, y::TrackedVar) . What if I wanted to implement the same idea, but which works with a general function? For example, z = f(x, y) . Any idea how to do that?

Typically, the way to do this using dispatch is to loop over all function names you want to overload and then use eval to programmatically create those method definitions, similar to this:

for f in [:(Base.:+), :f, ...]
    @eval begin
        function $f(x::TrackedVar, y::TrackedVar)
            return TrackedVar($f(x.value, y.value), [x, y])
        end
    end
end
1 Like

Hmmm what if some of those are functions with one argument, or three? :slight_smile: Or some of them have one argument of type TrackedVar, and second argument Float64?

You can just change the signature from $f(x::TrackedVar, y::TrackedVar) to $f(::TrackedVar), $f(::TrackedVar, ::Float64), etc. Although, for number types, it’s probably better to hook into the promotion mechanism:

1 Like

Thank you :slight_smile:

You are welcome!