move location utils here for now

Author: MikeInnes

src/MacroTools.jl

@@ -17,6 +17,7 @@ include("examples/threading.jl")
 include("examples/forward.jl")
 
 include("patch/diff.jl")
+include("patch/location.jl")
 include("patch/cst.jl")
 
 const animals = Symbol[]

src/patch/cst.jl

@@ -1,11 +1,11 @@
 using CSTParser
-using CSTParser: EXPR, Call, Location, expr_location, charrange
+using CSTParser: EXPR, Call
 
 expridx(x, ii) = (@assert isempty(ii); x)
 expridx(x::Expr, ii) = isempty(ii) ? x : expridx(x.args[ii[1]], ii[2:end])
 
 function precedence_level(cst::EXPR, loc::Location)
-  parent = cst[CSTParser.parent(loc)]
+  parent = cst[MacroTools.parent(loc)]
   if parent isa Union{CSTParser.BinaryOpCall,CSTParser.UnaryOpCall,CSTParser.ChainOpCall}
     Base.operator_precedence(Expr(parent).args[1])
   elseif parent isa Union{CSTParser.BinarySyntaxOpCall,CSTParser.UnarySyntaxOpCall}
@@ -18,29 +18,29 @@ end
 # Get the range at loc, including trailing trivia from the previous node.
 function separator_range(cst, loc)
   i = loc.ii[end]
-  loc = CSTParser.parent(loc)
-  start = charrange(cst, CSTParser.child(loc, i-1))[1][1]+1
-  stop = charrange(cst, CSTParser.child(loc, i))[1][end]
+  loc = MacroTools.parent(loc)
+  start = charrange(cst, child(loc, i-1))[1][1]+1
+  stop = charrange(cst, child(loc, i))[1][end]
   return start:stop
 end
 
 function separator(cst, loc, x::EXPR{Call}, i)
   length(x.args) == 3 && return ""
   length(x.args) == 4 && return ", "
-  separator_range(cst, CSTParser.child(loc, max(i-1,4)))
+  separator_range(cst, child(loc, max(i-1,4)))
 end
 
 function separator(cst, loc, x::EXPR{CSTParser.Block}, i)
-  out, in = charrange(cst, CSTParser.child(loc, max(i-1,1)))
+  out, in = charrange(cst, child(loc, max(i-1,1)))
   in[end]+1:out[end]
 end
 
 function separator(cst, loc, x::CSTParser.BinaryOpCall, i)
-  separator_range(cst, CSTParser.child(loc, 2))
+  separator_range(cst, child(loc, 2))
 end
 
 function separator(cst::EXPR, loc::Location)
-  parent = CSTParser.parent(loc)
+  parent = MacroTools.parent(loc)
   separator(cst, parent, cst[parent], loc.ii[end])
 end
 
@@ -68,7 +68,7 @@ function replacement(src::SourceFile, p::Insert)
   append && (p.idx[end] -= 1)
   loc = expr_location(src.cst, p.idx)
   # TODO handle cases like this more generally
-  src.cst[CSTParser.parent(loc)] isa EXPR{Call} && (loc.ii[end] = max(loc.ii[end], 2))
+  src.cst[MacroTools.parent(loc)] isa EXPR{Call} && (loc.ii[end] = max(loc.ii[end], 2))
   _, span = charrange(src.cst, loc)
   point = append ? span[end] : span[1]-1
   sep = separator(src.cst, loc)

src/patch/location.jl

@@ -0,0 +1,71 @@
+using CSTParser
+import CSTParser: EXPR, AbstractEXPR
+
+struct Location
+  ii::Vector{Int}
+end
+
+parent(l::Location) = Location(l.ii[1:end-1])
+child(l::Location, i) = Location([l.ii..., i])
+
+Base.getindex(ex::EXPR, i::Integer) = ex.args[i]
+Base.getindex(ex::AbstractEXPR, i) = collect(ex)[i]
+
+function Base.getindex(ex::AbstractEXPR, l::Location)
+  for i in l.ii
+    ex = ex[i]
+  end
+  return ex
+end
+
+function charrange(ex::AbstractEXPR, l::Location)
+    o = 0
+    for i in l.ii
+      for j = 1:i-1
+        o += ex[j].fullspan
+      end
+      ex = ex[i]
+    end
+    full = o .+ (1:ex.fullspan)
+    inner = o .+ ex.span
+    return full, inner
+end
+
+function expr_child_index(x::EXPR, n)
+  for (i, a) in enumerate(x.args)
+    a isa CSTParser.PUNCTUATION && continue
+    n -= 1
+    n == 0 && return i
+  end
+end
+
+function expr_child_index(x::EXPR{CSTParser.Const}, i)
+  @assert i == 1
+  return 2
+end
+
+expr_child_index(x::CSTParser.UnaryOpCall, i) = i
+expr_child_index(x::CSTParser.UnarySyntaxOpCall, _) = x.arg1 isa OPERATOR ? 2 : 1
+expr_child_index(x::CSTParser.BinaryOpCall, i) = [2, 1, 3][i]
+expr_child_index(x::CSTParser.BinarySyntaxOpCall, i) = i+1
+expr_child_index(x::CSTParser.ConditionalOpCall, i) = i
+expr_child_index(x::EXPR{CSTParser.ChainOpCall}, i) = i == 1 ? 2 : 2(i-1)-1
+
+function expr_child!(ii, ex, i)
+  j = expr_child_index(ex, i)
+  push!(ii, j)
+  return ex[j]
+end
+
+function expr_child!(ii, x::EXPR{<:Union{CSTParser.Begin,CSTParser.InvisBrackets}}, i)
+  push!(ii, 2, i)
+  return x[2][i]
+end
+
+function expr_location(ex::AbstractEXPR, ii)
+  jj = []
+  for i in ii
+    ex = expr_child!(jj, ex, i)
+  end
+  return Location(jj)
+end