Expand fn

3 files changed, 36 insertions, 5 deletions

diff --git a/lib/elixir/src/elixir_exp.erl b/lib/elixir/src/elixir_exp.erl
index 770d523cb..5d174ed62 100644
--- a/lib/elixir/src/elixir_exp.erl
+++ b/lib/elixir/src/elixir_exp.erl
@@ -1,5 +1,5 @@
 -module(elixir_exp).
--export([expand/2]).
+-export([expand/2, expand_many/2]).
 -import(elixir_errors, [compile_error/3, compile_error/4]).
 -include("elixir.hrl").
 
@@ -206,7 +206,6 @@ expand({ quote, Meta, [_, _] }, E) ->
 %% TODO: Remove me. Temporary during refactoring.
 expand({ '&', _, [Arg] } = Original, E) when is_integer(Arg) ->
   { Original, E };
-
 expand({ '&', Meta, [Arg] }, E) ->
   % assert_no_match_or_guard_scope(Meta, '&', S),
   case elixir_fn:capture(Meta, Arg, E) of
@@ -216,6 +215,10 @@ expand({ '&', Meta, [Arg] }, E) ->
       expand(Expr, EE)
   end;
 
+expand({ fn, Meta, Pairs }, E) ->
+  % assert_no_match_or_guard_scope(Meta, 'fn', S),
+  elixir_fn:expand(Meta, Pairs, E);
+
 %% Comprehensions
 
 expand({ Kind, Meta, Args }, E) when is_list(Args), (Kind == lc) orelse (Kind == bc) ->
diff --git a/lib/elixir/src/elixir_exp_clauses.erl b/lib/elixir/src/elixir_exp_clauses.erl
index e8ef50774..d923782be 100644
--- a/lib/elixir/src/elixir_exp_clauses.erl
+++ b/lib/elixir/src/elixir_exp_clauses.erl
@@ -2,9 +2,29 @@
 %% fn, receive and friends. try is handled in elixir_try.
 -module(elixir_exp_clauses).
 -export([match/3]).
+-export([expand_clause/3]).
 -include("elixir.hrl").
 
 match(Fun, Expr, #elixir_env{context=Context} = E) ->
   { EExpr, EE } = Fun(Expr, E#elixir_env{context=match}),
   { EExpr, EE#elixir_env{context=Context} }.
 
+expand_clause(Fun, { '->', Meta, [Left, Right] }, E) ->
+  { ELeft, EL }  = expand_head(Fun, Left, E),
+  { ERight, ER } = elixir_exp:expand(Right, EL),
+  { { '->', Meta, [ELeft, ERight] }, ER }.
+
+expand_head(Fun, [{ 'when', Meta, [_,_|_] = All }], E) ->
+  { Args, Guard } = elixir_utils:split_last(All),
+  { EArgs, EA }   = match(Fun, Args, E),
+  { EGuard, EG }  = expand_guard(Guard, EA#elixir_env{context=guard}),
+  { [{ 'when', Meta, EArgs ++ [EGuard] }], EG#elixir_env{context=E#elixir_env.context} };
+expand_head(Fun, Args, E) ->
+  match(Fun, Args, E).
+
+expand_guard({ 'when', Meta, [Left, Right] }, E) ->
+  { ELeft, EL }  = expand_guard(Left, E),
+  { ERight, ER } = expand_guard(Right, EL),
+  { { 'when', Meta, [ELeft, ERight] }, ER };
+expand_guard(Other, E) ->
+  elixir_exp:expand(Other, E).
diff --git a/lib/elixir/src/elixir_fn.erl b/lib/elixir/src/elixir_fn.erl
index f48642a89..20b6cc320 100644
--- a/lib/elixir/src/elixir_fn.erl
+++ b/lib/elixir/src/elixir_fn.erl
@@ -1,5 +1,5 @@
 -module(elixir_fn).
--export([fn/3, capture/3]).
+-export([fn/3, capture/3, expand/3]).
 -import(elixir_scope, [umergec/2]).
 -import(elixir_errors, [syntax_error/3, compile_error/4]).
 -include("elixir.hrl").
@@ -25,6 +25,16 @@ translate_fn_match(Arg, S) ->
   { TArg, TS } = elixir_translator:translate(Arg, S#elixir_scope{extra=fn_match}),
   { TArg, TS#elixir_scope{extra=S#elixir_scope.extra} }.
 
+%% Expansion
+
+expand(Meta, Clauses, E) ->
+  Transformer = fun(Clause, Acc) ->
+    { EClause, EC } = elixir_exp_clauses:expand_clause(fun elixir_exp:expand_many/2, Clause, Acc),
+    { EClause, elixir_env:mergec(E, EC) }
+  end,
+  { EClauses, _ } = lists:mapfoldl(Transformer, E, Clauses),
+  { { fn, Meta, EClauses }, E }.
+
 %% Capture
 
 capture(Meta, { '/', _, [{ { '.', _, [_, F] } = Dot, RequireMeta , [] }, A] }, E) when is_atom(F), is_integer(A) ->
@@ -72,8 +82,6 @@ capture(Meta, Arg, E) when is_integer(Arg) ->
 capture(Meta, Arg, E) ->
   invalid_capture(Meta, Arg, E).
 
-%% Helpers
-
 capture_import(Meta, { Atom, ImportMeta, Args } = Expr, E, Sequential) ->
   Res = Sequential andalso
         elixir_exp_dispatch:import_function(ImportMeta, Atom, length(Args), E),