Simple loop-construction macro for Julia

loop_ts.jl

using Base.Meta

macro loop_ts(ex)
  l, r = ex.args
  
  # Determine which symbol should become the loop index. We've assumed that
  # the equation has the form `y[t+1] = ...` or `y[t] = ...`. In the first case
  # we need to parse 't+1' to find the symbol (t); in the second case, we can
  # just read it off.
  idx =
    if isexpr(l.args[2], :call)
      filter(x -> isa(x, Symbol), l.args[2].args[2:end])[1]
    elseif isa(l.args[2], Symbol)
      l.args[2]
    end

  # Determine the start and end values for the loop index. The start and end
  # values are determined by parsing several terms like `y[t-1]`, `y[t+1]`,
  # `y[t-3]`, `e[t+1]`, etc. and pulling out `-1`, `+1`, `-3`, `+1`, then using
  # those integers to shift the endpoints from 1:length(y). `y` is also
  # determined automatically.
  offsets = extrema(vcat(get_offsets(l), get_offsets(r)))
  loopindex = :($(1 - offsets[1]):(length($(l.args[1])) - $(offsets[2])))

  quote
    for $idx in $loopindex
      $ex
    end
  end
end

# `get_offsets` expects to be given two types of expressions. The first
# is a function call representing `y[t+2] + 0.2y[t-3] + 0.5y[t]`; in
# that case we want to call `get_offsets` on each term individually.
# The second is an indvidual term `y[t-3]`; then we want to grab and parse the
# sybmolic array index `t-3`, which is done by the function `get_offset_from_rec`.
function get_offsets(ex_::Expr)
  isexpr(ex_,:call) && return [[get_offsets(a) for a in ex_.args[2:end]]...]
  isexpr(ex_,:ref) && return get_offset_from_ref(ex_.args[2])
  warning("Not expecting to be here")
  return Int64[]
end

get_offsets(x) = Int64[]

get_offset_from_ref(s::Symbol) = 0
get_offset_from_ref(x::Number) = x

function get_offset_from_ref(ex_::Expr)
  if isexpr(ex_,:call)
    ex_.args[1] == :+ &&
      return sum([get_offset_from_ref(a) for a in ex_.args[2:end]])
    ex_.args[1] == :- &&
      return (get_offset_from_ref(ex_.args[2])
              - sum([get_offset_from_ref(a) for a in ex_.args[3:end]]))
  end
  warning("Didn't expect to get here")
  return(0)
end