Scope question

[yebai]

Is anyone aware a package/way for extracting scoping structure of a Julia function? For a special reason, I need to implement a macro that can rename variables in child scopes so that variables that exist in both parent scope and child scopes are renamed to different names.

@mohamed82008 @ChrisRackauckas @MikeInnes @StefanKarpinski

[mohamed82008]

Why not just replace every symbol in the body of every function by a unique gensym() keeping track of the mapping of which symbol became which other symbol? This should be possible with @capture and postwalk of MacroTools.jl. To make the global-only variables visible and not error, you will need to replace the variables on the LHS of an assignment first. Then whatever was replaced from the LHS, you then replace those all over the function in another postwalk. To keep track of what was replaced you can do something like:

julia> using MacroTools

julia> mapping = Dict{Symbol, Symbol}()
Dict{Symbol,Symbol} with 0 entries

julia> MacroTools.postwalk(x -> x == :x ? (newx = gensym(x); mapping[x] = newx; newx) : x, :(x = 1))
:(##x#358 = 1)

julia> mapping
Dict{Symbol,Symbol} with 1 entry:
  :x => Symbol("##x#358")

This is one approach anyways.

[yebai]

Thanks, @mohamed82008. This problem becomes complicated when we re-use variable names between child and parent scopes. Consider the following example

using Turing

@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  z  ~ Bernoulli(0.5)
  x  ~ Normal(m+1, sqrt(s))  # x1

  let 
    if z == 1
        x ~ Normal(m+1.5, sqrt(s)) # x2
    else 
        x ~ Normal(m+2, sqrt(s))  # x2
    end
  end

  for i=1:5
     x ~ Normal(m+3, sqrt(s)) # x3
  end

  return s, m
end

Here we re-used variable names x for 3 different purposes. We want to rename x consistently so that x in different scopes are mapped to different names (e.g. x1, x2, x3). It's possible to solve this task using MacroTools, but I think it'll involve writing a scope parser, which might take a lot of time.

Related: #484

[mohamed82008]

This seems to work:

using MacroTools: @capture

function replace_vars(expr)
	mapping = Dict{Symbol, Symbol}()
	_find_vars(expr, mapping)
	_replace_vars(expr, mapping)
	expr
end
function _find_vars(expr, mapping)
	if expr isa Expr
		for i in 1:length(expr.args)
			x = expr.args[i]
			if (@capture(x, (V_ = R_)) || @capture(x, (V_ ~ R_))) && V isa Symbol
				if !haskey(mapping, V)
					newvar = gensym(V)
					mapping[V] = newvar
				end
				_find_vars(R, mapping)
			elseif @capture(x, for V_ in R_ body_ end | for V_ = R_ body_ end) || 
					@capture(x, let V_ = R_; body_ end | let V_; body_ end) || 
					@capture(x, f_(args__) = body_ | function f_(args__) body_ end)

				replace_vars(x)
			else
				_find_vars(x, mapping)
			end
		end
	end
	expr
end
function _replace_vars(expr, mapping)
	if expr isa Expr
		for i in 1:length(expr.args)
			x = expr.args[i]
			if haskey(mapping, x)
				expr.args[i] = mapping[x]
			else
				_replace_vars(x, mapping)
			end
		end
	end
	expr
end

julia> expr = :(@model gdemo(x) = begin
         s ~ InverseGamma(2,3)
         m ~ Normal(0,sqrt(s))
         z  ~ Bernoulli(0.5)
         x  ~ Normal(m+1, sqrt(s))  # x1

         let
           if z == 1
               x ~ Normal(m+1.5, sqrt(s)) # x2
           else
               x ~ Normal(m+2, sqrt(s))  # x2
           end
         end

         for i=1:5
            x ~ Normal(m+3, sqrt(s)) # x3
         end

         return s, m
       end);
julia> replace_vars(expr)
:(#= REPL[5]:1 =# @model gdemo(##x#361) = begin
              #= REPL[5]:1 =#
              #= REPL[5]:2 =#
              ##s#358 ~ InverseGamma(2, 3)
              #= REPL[5]:3 =#
              ##m#359 ~ Normal(0, sqrt(##s#358))
              #= REPL[5]:4 =#
              ##z#360 ~ Bernoulli(0.5)
              #= REPL[5]:5 =#
              ##x#361 ~ Normal(##m#359 + 1, sqrt(##s#358))
              #= REPL[5]:7 =#
              let
                  #= REPL[5]:8 =#
                  if ##z#360 == 1
                      #= REPL[5]:9 =#
                      ##x#362 ~ Normal(##m#359 + 1.5, sqrt(##s#358))
                  else
                      #= REPL[5]:11 =#
                      ##x#362 ~ Normal(##m#359 + 2, sqrt(##s#358))
                  end
              end
              #= REPL[5]:15 =#
              for ##i#363 = 1:5
                  #= REPL[5]:16 =#
                  ##x#364 ~ Normal(##m#359 + 3, sqrt(##s#358))
              end
              #= REPL[5]:19 =#
              return (##s#358, ##m#359)
          end)

The second branch of the if statement of _find_vars enumerates the new scope patterns. If you want to consider while loops, they will need to be added, or other function syntax, etc. This may not be a great approach but it gets the job done for the supported syntax.

[yebai]

Hi @mohamed82008, could you:

• turn this into a PR with some tests
• add support for multi-level of nesting scopes (e.g. let let end end)
• add support for other Julia language structs which can introduce new scopes.

Related: Forbid variables with same name #146