From 5bfc3336ca71af00c2330644dc3f7f96daa53f00 Mon Sep 17 00:00:00 2001
From: jkrumbiegel <22495855+jkrumbiegel@users.noreply.github.com>
Date: Wed, 8 Dec 2021 10:45:46 +0100
Subject: [PATCH] fix nested single line chain (#43)

* fix nested single line chain

* bump patch version
---
 Project.toml     | 2 +-
 src/Chain.jl     | 4 ++--
 test/runtests.jl | 6 ++++++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 5cee706..8f676d3 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Chain"
 uuid = "8be319e6-bccf-4806-a6f7-6fae938471bc"
 authors = ["Julius Krumbiegel"]
-version = "0.4.9"
+version = "0.4.10"
 
 [compat]
 julia = "1"
diff --git a/src/Chain.jl b/src/Chain.jl
index cc06f82..be144fa 100644
--- a/src/Chain.jl
+++ b/src/Chain.jl
@@ -257,7 +257,7 @@ function replace_underscores(expr::Expr, replacement)
     # if a @chain macrocall is found, only its first arg can be replaced if it's an
     # underscore, otherwise the macro insides are left untouched
     if expr.head == :macrocall && expr.args[1] == Symbol("@chain")
-        length(expr.args) != 4 && error("Malformed nested @chain macro")
+        length(expr.args) < 3 && error("Malformed nested @chain macro")
         expr.args[2] isa LineNumberNode || error("Malformed nested @chain macro")
         arg3 = if expr.args[3] == Symbol("_")
             found_underscore = true
@@ -265,7 +265,7 @@ function replace_underscores(expr::Expr, replacement)
         else
             expr.args[3]
         end
-        newexpr = Expr(:macrocall, Symbol("@chain"), expr.args[2], arg3, expr.args[4])
+        newexpr = Expr(:macrocall, Symbol("@chain"), expr.args[2], arg3, expr.args[4:end]...)
     # for all other expressions, their arguments are checked for underscores recursively
     # and replaced if any are found
     else
diff --git a/test/runtests.jl b/test/runtests.jl
index 6bfdb52..a5a2542 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -435,4 +435,10 @@ end
         "$_"
         uppercase
     end
+end
+
+@testset "nested single line chain" begin
+    @test 36 == @chain 1:3 begin
+        @chain _ sum _ ^ 2
+    end
 end
\ No newline at end of file