Support try/receive in record rewriter

Author: José Valim

lib/elixir/lib/kernel/record_rewriter.ex

@@ -104,6 +104,7 @@ defmodule Kernel.RecordRewriter do
 
   defp optimize_expr({ :tuple, line, args }, dict) do
     { args, dict, args_res } = optimize_tuple_args(args, dict)
+    args_res = if Enum.any?(args_res), do: args_res, else: nil
 
     res =
       case args do
@@ -116,7 +117,7 @@ defmodule Kernel.RecordRewriter do
 
   defp optimize_expr({ :var, _, name } = var, dict) do
     case :orddict.find(name, dict) do
-      { :ok, res } -> { var, dict, { res, nil } }
+      { :ok, res } -> { var, dict, res }
       :error       -> { var, dict, nil }
     end
   end
@@ -130,6 +131,38 @@ defmodule Kernel.RecordRewriter do
     { { :case, line, expr, clauses }, dict, res }
   end
 
+  defp optimize_expr({ :receive, line, clauses }, dict) do
+    tuples  = lc clause inlist clauses, do: optimize_clause(clause, dict)
+    clauses = lc { clause, _, _ } inlist tuples, do: clause
+    dict    = join_dict(tuples)
+    res     = join_result(tuples)
+    { { :receive, line, clauses }, dict, res }
+  end
+
+  defp optimize_expr({ :receive, line, clauses, after_key, after_value }, dict) do
+    tuples  = lc clause inlist clauses, do: optimize_clause(clause, dict)
+    clauses = lc { clause, _, _ } inlist tuples, do: clause
+
+    { after_key, dict, _ } = optimize_expr(after_key, dict)
+    { after_value, dict, res } = optimize_body(after_value, dict, [])
+
+    dict = join_dict(tuples, dict)
+    res  = join_result(tuples, res)
+
+    { { :receive, line, clauses, after_key, after_value }, dict, res }
+  end
+
+  defp optimize_expr({ :try, line, body, [], clauses, try_after }, dict) do
+    tuples  = lc clause inlist clauses, do: optimize_clause(clause, dict)
+    clauses = lc { clause, _, _ } inlist tuples, do: clause
+
+    { body, _, res } = optimize_body(body, dict, [])
+    res = join_result(tuples, res)
+
+    { try_after, _, _ } = optimize_body(try_after, dict, [])
+    { { :try, line, body, [], clauses, try_after }, dict, res }
+  end
+
   defp optimize_expr({ :fun, line, { :function, module, name, arity } }, dict) do
     { module, dict, _ } = optimize_expr(module, dict)
     { name, dict, _ }   = optimize_expr(name, dict)
@@ -179,14 +212,16 @@ defmodule Kernel.RecordRewriter do
     if is_record?(value) do
       dict =
         case :orddict.find(key, dict) do
-          { :ok, ^value } ->
+          { :ok, ^res } ->
             dict
+          { :ok, { ^value, _ } } ->
+            :orddict.store(key, { value, nil }, dict)
           { :ok, _ } ->
             # We are overriding a type of an existing variable,
             # which means the source code is invalid.
             :orddict.store(key, nil, dict)
           :error ->
-            :orddict.store(key, value, dict)
+            :orddict.store(key, res, dict)
         end
     end
 
@@ -232,9 +267,10 @@ defmodule Kernel.RecordRewriter do
 
   defp join_dict([{ _, dict, _ }|t], other) do
     other = Enum.reduce other, other, fn
-      { key, value }, acc ->
+      { key, { value, _ } = res }, acc ->
         case :orddict.find(key, dict) do
-          { :ok, ^value } -> acc
+          { :ok, ^res } -> acc
+          { :ok, { ^value, _ } } -> :orddict.store(key, { value, nil }, acc)
           { :ok, _ } -> :orddict.store(key, nil, acc)
           :error -> :orddict.erase(key, acc)
         end
@@ -255,6 +291,10 @@ defmodule Kernel.RecordRewriter do
     join_result(t, res)
   end
 
+  defp join_result([{ _, _, { res, _ } }|t], { res, _ }) do
+    join_result(t, { res, nil })
+  end
+
   defp join_result([{ _, _, _ }|_], _res) do
     nil
   end

lib/elixir/test/elixir/kernel/record_rewriter_test.exs

@@ -3,8 +3,6 @@ Code.require_file "../../test_helper.exs", __FILE__
 defmodule Kernel.RecordRewriterTest do
   use ExUnit.Case, async: true
 
-  import Kernel.RecordRewriter
-
   ## Helpers
 
   defmacrop clause(expr) do
@@ -17,6 +15,12 @@ defmodule Kernel.RecordRewriterTest do
     hd(clauses)
   end
 
+  defp optimize_clause(clause) do
+    { clause, dict, res } = Kernel.RecordRewriter.optimize_clause(clause)
+    dict = Enum.map dict, fn { k, { v, _ } } -> { k, v }; other -> other; end
+    { clause, dict, res }
+  end
+
   ## Dictionary tests
 
   test "simple atom" do
@@ -62,16 +66,21 @@ defmodule Kernel.RecordRewriterTest do
     assert optimize_clause(clause) == { clause, [x: Macro.Env, y: Range], nil }
   end
 
-  test "with body variable" do
+  test "inside body" do
     clause = clause(fn(x = Macro.Env[]) -> y = x end)
     assert optimize_clause(clause) == { clause, [x: Macro.Env, y: Macro.Env], { Macro.Env, nil } }
   end
 
-  test "with body variable overridden" do
+  test "inside body with variable overridden" do
     clause = clause(fn(x = Macro.Env[], y = Range[]) -> y = x end)
     assert optimize_clause(clause) == { clause, [x: Macro.Env, y: Range, "y@1": Macro.Env], { Macro.Env, nil } }
   end
 
+  test "inside body with nested tuple" do
+    clause = clause(fn(x = Range[]) -> ^x = Range[first: { :hello, 2 }] end)
+    assert optimize_clause(clause) == { clause, [x: Range], { Range, [nil, { :hello, nil }, nil] } }
+  end
+
   test "conflicting definition" do
     clause = clause(fn(x = Macro.Env[]) -> ^x = Range[]; :foo end)
     assert optimize_clause(clause) == { clause, [x: nil], nil }
@@ -148,4 +157,45 @@ defmodule Kernel.RecordRewriterTest do
     clause = clause(fn -> case something do 1 -> x = Macro.Env[]; 2 -> x = Range[] end end)
     assert optimize_clause(clause) == { clause, [x: nil], nil }
   end
+
+  test "inside case with nested tuple" do
+    clause = clause(fn -> case something do x = Range[first: { :foo, 2 }] -> x; Range[] = x -> x end end)
+    assert optimize_clause(clause) == { clause, [x: Range], { Range, nil } }
+
+    clause = clause(fn -> case something do x = Range[first: { :foo, 2 }] -> x; Range[first: { :foo, 2 }] = x -> x end end)
+    assert optimize_clause(clause) == { clause, [x: Range], { Range, [nil, { :foo, nil }, nil] } }
+  end
+
+  test "inside receive" do
+    clause = clause(fn -> receive do x = Macro.Env[] -> x; Macro.Env[] = x -> x end end)
+    assert optimize_clause(clause) == { clause, [x: Macro.Env], { Macro.Env, nil } }
+
+    clause = clause(fn -> receive do x = Macro.Env[] -> x; Macro.Env[] = y -> y end end)
+    assert optimize_clause(clause) == { clause, [], { Macro.Env, nil } }
+
+    clause = clause(fn -> receive do 1 -> x = Macro.Env[]; 2 -> x = Macro.Env[] end end)
+    assert optimize_clause(clause) == { clause, [x: Macro.Env], { Macro.Env, nil } }
+
+    clause = clause(fn -> receive do 1 -> x = Macro.Env[]; 2 -> x = Range[] end end)
+    assert optimize_clause(clause) == { clause, [x: nil], nil }
+  end
+
+  test "inside receive with after" do
+    clause = clause(fn -> receive do 1 -> x = Macro.Env[]; after 2 -> x = Macro.Env[]; end end)
+    assert optimize_clause(clause) == { clause, [x: Macro.Env], { Macro.Env, nil } }
+
+    clause = clause(fn -> receive do 1 -> x = Macro.Env[]; after 2 -> x = Range[]; end end)
+    assert optimize_clause(clause) == { clause, [x: nil], nil }
+  end
+
+  test "inside try" do
+    clause = clause(fn -> try do x = Macro.Env[]; x end end)
+    assert optimize_clause(clause) == { clause, [], { Macro.Env, nil } }
+
+    clause = clause(fn -> try do x = Macro.Env[]; x; catch :oops -> x = Macro.Env[]; end end)
+    assert optimize_clause(clause) == { clause, [], { Macro.Env, nil } }
+
+    clause = clause(fn -> try do x = Macro.Env[]; x; after x = Macro.Env[]; end end)
+    assert optimize_clause(clause) == { clause, [], { Macro.Env, nil } }
+  end
 end