Repository: domluna/JuliaFormatter.jl
Pull request #79 (2019-09-27 03:55:55)

Add nospace and nonest for ChainOpCall
 


Patch:
diff --git a/src/pretty.jl b/src/pretty.jl
index 15bec6b7..d0ba87f7 100644
--- a/src/pretty.jl
+++ b/src/pretty.jl
@@ -1052,14 +1052,14 @@ function p_if(x, s)
 end
 
 # ChainOpCall/Comparison
-function p_chaincall(x, s)
+function p_chaincall(x, s; nonest = false, nospace = false)
     t = PTree(x, nspaces(s))
     for (i, a) in enumerate(x)
         n = pretty(a, s)
         if a.typ === CSTParser.OPERATOR
-            add_node!(t, Whitespace(1), s)
+            !nospace && add_node!(t, Whitespace(1), s)
             add_node!(t, n, s, join_lines = true)
-            add_node!(t, Placeholder(1), s)
+            !(nospace && nonest) && add_node!(t, Placeholder(1), s)
         elseif i == length(x) - 1 && is_punc(a) && is_punc(x.args[i+1])
             add_node!(t, n, s, join_lines = true)
         elseif CSTParser.is_comma(a) && i != length(x)
@@ -1168,6 +1168,8 @@ function p_binarycall(x, s; nonest = false, nospace = false)
         add_node!(t, p_binarycall(x.args[1], s, nonest = nonest, nospace = nospace), s)
     elseif x.args[1].typ === CSTParser.InvisBrackets
         add_node!(t, p_invisbrackets(x.args[1], s, nonest = nonest, nospace = nospace), s)
+    elseif x.args[1].typ === CSTParser.ChainOpCall
+        add_node!(t, p_chaincall(x.args[1], s, nonest = nonest, nospace = nospace), s)
     else
         add_node!(t, pretty(x.args[1], s), s)
     end
@@ -1201,6 +1203,8 @@ function p_binarycall(x, s; nonest = false, nospace = false)
         n = p_binarycall(x.args[3], s, nonest = nonest, nospace = nospace)
     elseif x.args[3].typ === CSTParser.InvisBrackets
         n = p_invisbrackets(x.args[3], s, nonest = nonest, nospace = nospace)
+    elseif x.args[3].typ === CSTParser.ChainOpCall
+        n = p_chaincall(x.args[3], s, nonest = nonest, nospace = nospace)
     else
         n = pretty(x.args[3], s)
     end


Test patch:
diff --git a/test/files/reverse.jl b/test/files/reverse.jl
index ea1e34fa..90f1a5e5 100644
--- a/test/files/reverse.jl
+++ b/test/files/reverse.jl
@@ -39,8 +39,8 @@ xgradindex(x, i) = xcall(Zygote, :gradindex, x, i)
 normalise!(ir) = ir |> IRTools.merge_returns!
 
 function instrument_new!(ir, v, ex)
-  isexpr(ex, :new) ? (ir[v] = xcall(Zygote, :__new__, ex.args...)) :
-  isexpr(ex, :splatnew) ? (ir[v] = xcall(Zygote, :__splatnew__, ex.args...)) : ex
+    isexpr(ex, :new) ? (ir[v] = xcall(Zygote, :__new__, ex.args...)) :
+    isexpr(ex, :splatnew) ? (ir[v] = xcall(Zygote, :__splatnew__, ex.args...)) : ex
 end
 
 # Hack to work around fragile constant prop through overloaded functions
@@ -48,108 +48,118 @@ unwrapquote(x) = x
 unwrapquote(x::QuoteNode) = x.value
 
 is_literal_getproperty(ex) =
-  (iscall(ex, Base, :getproperty) ||
-   iscall(ex, Core, :getfield) || iscall(ex, Base, :getfield)) &&
-  ex.args[3] isa Union{QuoteNode,Integer}
+    (iscall(ex, Base, :getproperty) ||
+     iscall(ex, Core, :getfield) || iscall(ex, Base, :getfield)) &&
+    ex.args[3] isa Union{QuoteNode,Integer}
 
 function instrument_getproperty!(ir, v, ex)
-  is_literal_getproperty(ex) ?
-  (ir[v] = xcall(Zygote, :literal_getproperty, ex.args[2], Val(unwrapquote(ex.args[3])))) :
-  ex
+    is_literal_getproperty(ex) ?
+    (ir[v] = xcall(
+        Zygote,
+        :literal_getproperty,
+        ex.args[2],
+        Val(unwrapquote(ex.args[3])),
+    )) :
+    ex
 end
 
 is_literal_getindex(ex) =
-  iscall(ex, Base, :getindex) &&
-  length(ex.args) == 3 && ex.args[3] isa Union{Integer,QuoteNode}
+    iscall(ex, Base, :getindex) &&
+    length(ex.args) == 3 && ex.args[3] isa Union{Integer,QuoteNode}
 
 function instrument_getindex!(ir, v, ex)
-  is_literal_getindex(ex) ?
-  (ir[v] = xcall(Zygote, :literal_getindex, ex.args[2], Val(unwrapquote(ex.args[3])))) : ex
+    is_literal_getindex(ex) ?
+    (ir[v] = xcall(Zygote, :literal_getindex, ex.args[2], Val(unwrapquote(ex.args[3])))) :
+    ex
 end
 
 is_literal_iterate(ex) =
-  iscall(ex, Base, :indexed_iterate) &&
-  length(ex.args) >= 3 && ex.args[3] isa Union{Integer,QuoteNode}
+    iscall(ex, Base, :indexed_iterate) &&
+    length(ex.args) >= 3 && ex.args[3] isa Union{Integer,QuoteNode}
 
 function instrument_iterate!(ir, v, ex)
-  is_literal_iterate(ex) ?
-  (ir[v] = xcall(
-    Zygote,
-    :literal_indexed_iterate,
-    ex.args[2],
-    Val(unwrapquote(ex.args[3])),
-    ex.args[4:end]...,
-  )) :
-  ex
+    is_literal_iterate(ex) ?
+    (ir[v] = xcall(
+        Zygote,
+        :literal_indexed_iterate,
+        ex.args[2],
+        Val(unwrapquote(ex.args[3])),
+        ex.args[4:end]...,
+    )) :
+    ex
 end
 
 function instrument_literals!(ir, v, ex)
-  ex = instrument_getproperty!(ir, v, ex)
-  ex = instrument_getindex!(ir, v, ex)
-  ex = instrument_iterate!(ir, v, ex)
+    ex = instrument_getproperty!(ir, v, ex)
+    ex = instrument_getindex!(ir, v, ex)
+    ex = instrument_iterate!(ir, v, ex)
 end
 
 function istrackable(x)
-  x isa GlobalRef && x.mod ∉ (Base, Core) || return false
-  isconst(x.mod, x.name) || return true
-  x = getfield(x.mod, x.name)
-  !(x isa Type || sizeof(x) == 0)
+    x isa GlobalRef && x.mod ∉ (Base, Core) || return false
+    isconst(x.mod, x.name) || return true
+    x = getfield(x.mod, x.name)
+    !(x isa Type || sizeof(x) == 0)
 end
 
 function instrument_global!(ir, v, ex)
-  if istrackable(ex)
-    ir[v] = xcall(Zygote, :unwrap, QuoteNode(ex), ex)
-  else
-    ir[v] =
-      prewalk(ex) do x
-        istrackable(x) || return x
-        insert!(ir, v, stmt(xcall(Zygote, :unwrap, QuoteNode(x), x), type = exprtype(x)))
-      end
-  end
+    if istrackable(ex)
+        ir[v] = xcall(Zygote, :unwrap, QuoteNode(ex), ex)
+    else
+        ir[v] =
+            prewalk(ex) do x
+                istrackable(x) || return x
+                insert!(
+                    ir,
+                    v,
+                    stmt(xcall(Zygote, :unwrap, QuoteNode(x), x), type = exprtype(x)),
+                )
+            end
+    end
 end
 
 function instrument(ir::IR)
-  pr = Pipe(ir)
-  for (v, st) in pr
-    ex = st.expr
-    isexpr(ex, :foreigncall) && continue
-    isexpr(ex, :enter, :leave) && error("try/catch is not supported.")
-    ex = instrument_new!(pr, v, ex)
-    ex = instrument_literals!(pr, v, ex)
-    ex = instrument_global!(pr, v, ex)
-  end
-  return finish(pr)
+    pr = Pipe(ir)
+    for (v, st) in pr
+        ex = st.expr
+        isexpr(ex, :foreigncall) && continue
+        isexpr(ex, :enter, :leave) && error("try/catch is not supported.")
+        ex = instrument_new!(pr, v, ex)
+        ex = instrument_literals!(pr, v, ex)
+        ex = instrument_global!(pr, v, ex)
+    end
+    return finish(pr)
 end
 
 const BranchNumber = UInt8
 
 function record_branches!(ir::IR)
-  brs = Dict{Int,Variable}()
-  for bb in blocks(ir)
-    preds = predecessors(bb)
-    length(preds) > 1 || continue
-    brs[bb.id] = argument!(bb, BranchNumber(0), BranchNumber)
-    i = length(arguments(bb))
-    n = 0
-    for aa in blocks(ir), br in branches(aa)
-      br.block == bb.id && (arguments(br)[i] = BranchNumber(n += 1))
+    brs = Dict{Int,Variable}()
+    for bb in blocks(ir)
+        preds = predecessors(bb)
+        length(preds) > 1 || continue
+        brs[bb.id] = argument!(bb, BranchNumber(0), BranchNumber)
+        i = length(arguments(bb))
+        n = 0
+        for aa in blocks(ir), br in branches(aa)
+            br.block == bb.id && (arguments(br)[i] = BranchNumber(n += 1))
+        end
     end
-  end
-  return ir, brs
+    return ir, brs
 end
 
 ignored_f(f) =
-  f in (
-    GlobalRef(Base, :not_int),
-    GlobalRef(Core.Intrinsics, :not_int),
-    GlobalRef(Core, :(===)),
-    GlobalRef(Core, :apply_type),
-    GlobalRef(Core, :typeof),
-    GlobalRef(Core, :throw),
-    GlobalRef(Base, :kwerr),
-    GlobalRef(Core, :kwfunc),
-    GlobalRef(Core, :isdefined),
-  )
+    f in (
+        GlobalRef(Base, :not_int),
+        GlobalRef(Core.Intrinsics, :not_int),
+        GlobalRef(Core, :(===)),
+        GlobalRef(Core, :apply_type),
+        GlobalRef(Core, :typeof),
+        GlobalRef(Core, :throw),
+        GlobalRef(Base, :kwerr),
+        GlobalRef(Core, :kwfunc),
+        GlobalRef(Core, :isdefined),
+    )
 ignored_f(ir, f) = ignored_f(f)
 ignored_f(ir, f::Variable) = ignored_f(get(ir, f, nothing))
 
@@ -158,73 +168,73 @@ ignored(ir, ex::Variable) = ignored(ir, ir[ex])
 
 # TODO: remove this once we don't mess with type inference
 function _forward_type(Ts)
-  usetyped || return Any
-  all(T -> isconcretetype(T) || T <: DataType, Ts) || return Any
-  T = Core.Compiler.return_type(_forward, Tuple{Context,Ts...})
-  return T == Union{} ? Any : T
+    usetyped || return Any
+    all(T -> isconcretetype(T) || T <: DataType, Ts) || return Any
+    T = Core.Compiler.return_type(_forward, Tuple{Context,Ts...})
+    return T == Union{} ? Any : T
 end
 
 isvalidtype(jT, yT) = jT <: Tuple && length(jT.parameters) == 2 && jT.parameters[1] <: yT
 
 function primal(ir::IR)
-  pr = Pipe(ir)
-  pbs = Dict{Variable,Variable}()
-  argument!(pr, at = 1)
-  cx = argument!(pr, Context, at = 2)
-  for (v, st) in pr
-    ex = st.expr
-    if isexpr(ex, :call) && !ignored(ir, ex)
-      yT = exprtype(ir, v)
-      T = _forward_type(exprtype.((ir,), ex.args))
-      if yT == Any || isvalidtype(T, yT)
-        yJ = insert!(
-          pr,
-          v,
-          stmt(xcall(Zygote, :_forward, cx, ex.args...), line = ir[v].line),
-        )
-        pr[v] = xgetindex(yJ, 1)
-        J = insertafter!(
-          pr,
-          v,
-          stmt(
-            xgetindex(yJ, 2),
-            type = T == Any ? Any : T.parameters[2],
-            line = ir[v].line,
-          ),
-        )
-        pbs[v] = substitute(pr, J)
-      else
-        yJ = insert!(pr, v, xcall(Zygote, :_forward, cx, ex.args...))
-        y = insert!(pr, v, xgetindex(yJ, 1))
-        J = insert!(pr, v, stmt(xgetindex(yJ, 2), line = ir[v].line))
-        pr[v] = xcall(Zygote, :typeassert, y, yT)
-        pbs[v] = substitute(pr, J)
-      end
+    pr = Pipe(ir)
+    pbs = Dict{Variable,Variable}()
+    argument!(pr, at = 1)
+    cx = argument!(pr, Context, at = 2)
+    for (v, st) in pr
+        ex = st.expr
+        if isexpr(ex, :call) && !ignored(ir, ex)
+            yT = exprtype(ir, v)
+            T = _forward_type(exprtype.((ir,), ex.args))
+            if yT == Any || isvalidtype(T, yT)
+                yJ = insert!(
+                    pr,
+                    v,
+                    stmt(xcall(Zygote, :_forward, cx, ex.args...), line = ir[v].line),
+                )
+                pr[v] = xgetindex(yJ, 1)
+                J = insertafter!(
+                    pr,
+                    v,
+                    stmt(
+                        xgetindex(yJ, 2),
+                        type = T == Any ? Any : T.parameters[2],
+                        line = ir[v].line,
+                    ),
+                )
+                pbs[v] = substitute(pr, J)
+            else
+                yJ = insert!(pr, v, xcall(Zygote, :_forward, cx, ex.args...))
+                y = insert!(pr, v, xgetindex(yJ, 1))
+                J = insert!(pr, v, stmt(xgetindex(yJ, 2), line = ir[v].line))
+                pr[v] = xcall(Zygote, :typeassert, y, yT)
+                pbs[v] = substitute(pr, J)
+            end
+        end
     end
-  end
-  pr = finish(pr)
-  pr, brs = record_branches!(pr)
-  return pr, brs, pbs
+    pr = finish(pr)
+    pr, brs = record_branches!(pr)
+    return pr, brs, pbs
 end
 
 struct Primal
-  ir::IR
-  pr::IR
-  varargs::Union{Int,Nothing}
-  branches::Dict{Int,Variable}
-  pullbacks::Dict{Variable,Variable}
+    ir::IR
+    pr::IR
+    varargs::Union{Int,Nothing}
+    branches::Dict{Int,Variable}
+    pullbacks::Dict{Variable,Variable}
 end
 
 function Primal(ir::IR; varargs = nothing)
-  ir = instrument(normalise!(ir))
-  pr, brs, pbs = primal(ir)
-  Primal(expand!(ir), pr, varargs, brs, pbs)
+    ir = instrument(normalise!(ir))
+    pr, brs, pbs = primal(ir)
+    Primal(expand!(ir), pr, varargs, brs, pbs)
 end
 
 # Backwards Pass
 
 struct Alpha
-  id::Int
+    id::Int
 end
 
 Base.show(io::IO, x::Alpha) = print(io, "@", x.id)
@@ -234,112 +244,115 @@ alpha(x::Variable) = Alpha(x.id)
 Variable(a::Alpha) = Variable(a.id)
 
 sig(b::IRTools.Block) =
-  unique([arg for br in branches(b) for arg in br.args if arg isa Variable])
+    unique([arg for br in branches(b) for arg in br.args if arg isa Variable])
 sig(pr::Primal) = Dict(b.id => sig(b) for b in blocks(pr.ir))
 
 # TODO unreachables?
 function adjointcfg(pr::Primal)
-  ir = empty(pr.ir)
-  return!(ir, nothing)
-  for b in blocks(pr.ir)[2:end]
-    block!(ir)
-    preds = predecessors(b)
-    rb = block(ir, b.id)
-    for i = 1:length(preds)
-      cond = i == length(preds) ? nothing :
-             push!(rb, xcall(Base, :(!==), alpha(pr.branches[b.id]), BranchNumber(i)))
-      branch!(rb, preds[i].id, unless = cond)
+    ir = empty(pr.ir)
+    return!(ir, nothing)
+    for b in blocks(pr.ir)[2:end]
+        block!(ir)
+        preds = predecessors(b)
+        rb = block(ir, b.id)
+        for i = 1:length(preds)
+            cond = i == length(preds) ? nothing :
+                   push!(rb, xcall(Base, :(!==), alpha(pr.branches[b.id]), BranchNumber(i)))
+            branch!(rb, preds[i].id, unless = cond)
+        end
+        if !isempty(branches(b)) && branches(b)[end] == IRTools.unreachable
+            branch!(rb, 0)
+        end
     end
-    if !isempty(branches(b)) && branches(b)[end] == IRTools.unreachable
-      branch!(rb, 0)
+    sigs = sig(pr)
+    for b in blocks(ir)[1:end-1], i = 1:length(sigs[b.id])
+        argument!(b)
     end
-  end
-  sigs = sig(pr)
-  for b in blocks(ir)[1:end-1], i = 1:length(sigs[b.id])
-    argument!(b)
-  end
-  argument!(blocks(ir)[end])
-  return ir, sigs
+    argument!(blocks(ir)[end])
+    return ir, sigs
 end
 
 branchfor(ir, (from, to)) =
-  get(filter(br -> br.block == to, branches(block(ir, from))), 1, nothing)
+    get(filter(br -> br.block == to, branches(block(ir, from))), 1, nothing)
 
 xaccum(ir) = nothing
 xaccum(ir, x) = x
 xaccum(ir, xs...) = push!(ir, xcall(Zygote, :accum, xs...))
 
 function adjoint(pr::Primal)
-  ir, sigs = adjointcfg(pr)
-  for b in reverse(blocks(pr.ir))
-    rb = block(ir, b.id)
-    grads = Dict()
-    grad(x, x̄) = push!(get!(grads, x, []), x̄)
-    grad(x) = xaccum(rb, get(grads, x, [])...)
+    ir, sigs = adjointcfg(pr)
+    for b in reverse(blocks(pr.ir))
+        rb = block(ir, b.id)
+        grads = Dict()
+        grad(x, x̄) = push!(get!(grads, x, []), x̄)
+        grad(x) = xaccum(rb, get(grads, x, [])...)
     # Backprop through (successor) branch arguments
-    for i = 1:length(sigs[b.id])
-      grad(sigs[b.id][i], arguments(rb)[i])
-    end
+        for i = 1:length(sigs[b.id])
+            grad(sigs[b.id][i], arguments(rb)[i])
+        end
     # Backprop through statements
-    for v in reverse(keys(b))
-      ex = b[v].expr
-      if haskey(pr.pullbacks, v)
-        g = push!(rb, stmt(Expr(:call, alpha(pr.pullbacks[v]), grad(v)), line = b[v].line))
-        for (i, x) in enumerate(ex.args)
-          x isa Variable || continue
-          grad(x, push!(rb, stmt(xgradindex(g, i), line = b[v].line)))
+        for v in reverse(keys(b))
+            ex = b[v].expr
+            if haskey(pr.pullbacks, v)
+                g = push!(
+                    rb,
+                    stmt(Expr(:call, alpha(pr.pullbacks[v]), grad(v)), line = b[v].line),
+                )
+                for (i, x) in enumerate(ex.args)
+                    x isa Variable || continue
+                    grad(x, push!(rb, stmt(xgradindex(g, i), line = b[v].line)))
+                end
+            elseif ex isa Core.PiNode
+                grads[ex.val] = grads[v]
+            elseif isexpr(ex, GlobalRef, :call, :isdefined, :inbounds, :meta)
+            elseif isexpr(ex)
+                push!(
+                    rb,
+                    stmt(
+                        xcall(Base, :error, "Can't differentiate $(ex.head) expression"),
+                        line = b[v].line,
+                    ),
+                )
+            else # A literal value
+                continue
+            end
         end
-      elseif ex isa Core.PiNode
-        grads[ex.val] = grads[v]
-      elseif isexpr(ex, GlobalRef, :call, :isdefined, :inbounds, :meta)
-      elseif isexpr(ex)
-        push!(
-          rb,
-          stmt(
-            xcall(Base, :error, "Can't differentiate $(ex.head) expression"),
-            line = b[v].line,
-          ),
-        )
-      else # A literal value
-        continue
-      end
-    end
-    if b.id > 1 # Backprop through (predecessor) branch arguments
-      gs = grad.(arguments(b))
-      for br in branches(rb)
-        br.block == 0 && continue
-        br′ = branchfor(pr.ir, br.block => b.id)
-        br′ == nothing && continue
-        ins = br′.args
-        for i = 1:length(br.args)
-          ā = [gs[j] for j = 1:length(ins) if ins[j] == sigs[br.block][i]]
-          br.args[i] = xaccum(rb, ā...)
+        if b.id > 1 # Backprop through (predecessor) branch arguments
+            gs = grad.(arguments(b))
+            for br in branches(rb)
+                br.block == 0 && continue
+                br′ = branchfor(pr.ir, br.block => b.id)
+                br′ == nothing && continue
+                ins = br′.args
+                for i = 1:length(br.args)
+                    ā = [gs[j] for j = 1:length(ins) if ins[j] == sigs[br.block][i]]
+                    br.args[i] = xaccum(rb, ā...)
+                end
+            end
+        else # Backprop function arguments
+            gs = [grad(arg) for arg in arguments(pr.ir)]
+            Δ = push!(
+                rb,
+                pr.varargs == nothing ? xcall(Zygote, :tuple, gs...) :
+                xcall(Zygote, :tuple_va, Val(pr.varargs), gs...),
+            )
+            branches(rb)[1].args[1] = Δ
         end
-      end
-    else # Backprop function arguments
-      gs = [grad(arg) for arg in arguments(pr.ir)]
-      Δ = push!(
-        rb,
-        pr.varargs == nothing ? xcall(Zygote, :tuple, gs...) :
-        xcall(Zygote, :tuple_va, Val(pr.varargs), gs...),
-      )
-      branches(rb)[1].args[1] = Δ
     end
-  end
-  return ir
+    return ir
 end
 
 struct Adjoint
-  primal::IR
-  adjoint::IR
+    primal::IR
+    adjoint::IR
 end
 
 function Adjoint(ir::IR; varargs = nothing, normalise = true)
-  pr = Primal(ir, varargs = varargs)
-  adj = adjoint(pr) |> prune!
-  if normalise
-    permute!(adj, length(adj.blocks):-1:1)
-    adj = IRTools.domorder!(adj) |> IRTools.renumber
-  end
-  Adjoint(pr.pr, adj)
+    pr = Primal(ir, varargs = varargs)
+    adj = adjoint(pr) |> prune!
+    if normalise
+        permute!(adj, length(adj.blocks):-1:1)
+        adj = IRTools.domorder!(adj) |> IRTools.renumber
+    end
+    Adjoint(pr.pr, adj)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 8c014fc0..58b1310d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -242,6 +242,14 @@ end
         @test fmt("2a") == "2a"
         @test fmt("2(a+1)") == "2 * (a + 1)"
 
+        str_ = "a[1:2 * num_source * num_dump-1]"
+        str = "a[1:2*num_source*num_dump-1]"
+        @test fmt(str_, 4, 1) == str
+
+        str_ = "a[2 * num_source * num_dump-1:1]"
+        str = "a[2*num_source*num_dump-1:1]"
+        @test fmt(str_, 4, 1) == str
+
         str = "!(typ <: ArithmeticTypes)"
         @test fmt(str) == str