From 1e6b9c8cab42164f7f4bb6e9f64da94637e9b5a3 Mon Sep 17 00:00:00 2001
From: leaf corcoran <leafot@gmail.com>
Date: Mon, 1 Oct 2012 12:30:16 -0700
Subject: [PATCH] allow assignment in elseif

---
 moonscript/parse.lua      |  6 ++++--
 moonscript/transform.lua  | 29 +++++++++++++++++++++++++--
 moonscript/transform.moon | 24 ++++++++++++++++++----
 tests/inputs/cond.moon    | 23 +++++++++++++++++++++
 tests/outputs/cond.lua    | 42 ++++++++++++++++++++++++++++++++++++++-
 5 files changed, 115 insertions(+), 9 deletions(-)

diff --git a/moonscript/parse.lua b/moonscript/parse.lua
index c1e06f8..79f978b 100644
--- a/moonscript/parse.lua
+++ b/moonscript/parse.lua
@@ -398,8 +398,10 @@ local build_grammar = wrap_env(function()
 		SwitchCase = key"when" * Exp * key"then"^-1 * Body / mark"case",
 		SwitchElse = key"else" * Body / mark"else",
 
-		If = key"if" * (Exp * Assign^-1 / format_assign_for_if) * key"then"^-1 * Body *
-			((Break * CheckIndent)^-1 * EmptyLine^0 * key"elseif" * Exp * key"then"^-1 * Body / mark"elseif")^0 *
+		IfCond = Exp * Assign^-1 / format_assign_for_if,
+
+		If = key"if" * IfCond * key"then"^-1 * Body *
+			((Break * CheckIndent)^-1 * EmptyLine^0 * key"elseif" * IfCond * key"then"^-1 * Body / mark"elseif")^0 *
 			((Break * CheckIndent)^-1 * EmptyLine^0 * key"else" * Body / mark"else")^-1 / mark"if",
 
 		While = key"while" * Exp * key"do"^-1 * Body / mark"while",
diff --git a/moonscript/transform.lua b/moonscript/transform.lua
index d01e8cd..644ceae 100644
--- a/moonscript/transform.lua
+++ b/moonscript/transform.lua
@@ -250,6 +250,30 @@ hoist_declarations = function(body, rules)
     assigns
   })
 end
+local expand_elseif_assign
+expand_elseif_assign = function(ifstm)
+  for i = 4, #ifstm do
+    local case = ifstm[i]
+    if ntype(case) == "elseif" and ntype(case[2]) == "assign" then
+      local split = {
+        unpack(ifstm, 1, i - 1)
+      }
+      insert(split, {
+        "else",
+        {
+          {
+            "if",
+            case[2],
+            case[3],
+            unpack(ifstm, i + 1)
+          }
+        }
+      })
+      return split
+    end
+  end
+  return ifstm
+end
 local constructor_name = "new"
 local Transformer
 Transformer = (function()
@@ -526,8 +550,7 @@ Statement = Transformer({
     return node
   end,
   ["if"] = function(self, node, ret)
-    smart_node(node)
-    if ntype(node.cond) == "assign" then
+    if ntype(node[2]) == "assign" then
       local _, assign, body = unpack(node)
       local name = assign[2][1]
       return build["do"]({
@@ -539,7 +562,9 @@ Statement = Transformer({
         }
       })
     end
+    node = expand_elseif_assign(node)
     if ret then
+      smart_node(node)
       node['then'] = apply_to_last(node['then'], ret)
       for i = 4, #node do
         local case = node[i]
diff --git a/moonscript/transform.moon b/moonscript/transform.moon
index d512e14..0171593 100644
--- a/moonscript/transform.moon
+++ b/moonscript/transform.moon
@@ -100,6 +100,20 @@ hoist_declarations = (body, rules) ->
 
   table.insert body, 1, {"declare", assigns}
 
+expand_elseif_assign = (ifstm) ->
+  for i = 4, #ifstm
+    case = ifstm[i]
+    if ntype(case) == "elseif" and ntype(case[2]) == "assign"
+      split = { unpack ifstm, 1, i - 1 }
+      insert split, {
+        "else", {
+          {"if", case[2], case[3], unpack ifstm, i + 1}
+        }
+      }
+      return split
+
+  ifstm
+
 constructor_name = "new"
 
 class Transformer
@@ -231,10 +245,9 @@ Statement = Transformer {
     node
 
   if: (node, ret) =>
-    smart_node node
 
-    -- extract assigns in cond
-    if ntype(node.cond) == "assign"
+    -- expand assign in cond
+    if ntype(node[2]) == "assign"
       _, assign, body = unpack node
       name = assign[2][1]
       return build["do"] {
@@ -242,8 +255,11 @@ Statement = Transformer {
         {"if", name, unpack node, 3}
       }
 
-    -- handle cascading return decorator
+    node = expand_elseif_assign node
+
+    -- apply cascading return decorator
     if ret
+      smart_node node
       -- mutate all the bodies
       node['then'] = apply_to_last node['then'], ret
       for i = 4, #node
diff --git a/tests/inputs/cond.moon b/tests/inputs/cond.moon
index 17b894e..6b548f6 100644
--- a/tests/inputs/cond.moon
+++ b/tests/inputs/cond.moon
@@ -66,3 +66,26 @@ else
 hello = 5 + if something = 10
   print something
 
+---
+
+z = false
+
+if false
+  one
+elseif x = true
+  two
+elseif z = true
+  three
+else
+  four
+
+
+out = if false
+  one
+elseif x = true
+  two
+elseif z = true
+  three
+else
+  four
+
diff --git a/tests/outputs/cond.lua b/tests/outputs/cond.lua
index 47ea285..8de55bd 100644
--- a/tests/outputs/cond.lua
+++ b/tests/outputs/cond.lua
@@ -109,4 +109,44 @@ hello = 5 + (function()
       return print(something)
     end
   end
-end)()
\ No newline at end of file
+end)()
+local z = false
+if false then
+  local _ = one
+else
+  do
+    local x = true
+    if x then
+      local _ = two
+    else
+      do
+        z = true
+        if z then
+          local _ = three
+        else
+          local _ = four
+        end
+      end
+    end
+  end
+end
+local out
+if false then
+  out = one
+else
+  do
+    local x = true
+    if x then
+      out = two
+    else
+      do
+        z = true
+        if z then
+          out = three
+        else
+          out = four
+        end
+      end
+    end
+  end
+end
\ No newline at end of file