Allow &() as a general capture mechanism

Author: José Valim

lib/elixir/lib/kernel/special_forms.ex

@@ -869,6 +869,57 @@ defmodule Kernel.SpecialForms do
 
       &local_function/1
 
+  A capture also allows the captured functions to be partially
+  applied, for example:
+
+      iex> fun = &atom_to_binary(&1, :utf8)
+      iex> fun.(:hello)
+      "hello"
+
+  In the example above, we use &1 as a placeholder, generating
+  a function with one argument. We could also use `&2` and `&3`
+  to delimit more arguments:
+
+      iex> fun = &atom_to_binary(&1, &2)
+      iex> fun.(:hello, :utf8)
+      "hello"
+
+  Since operators are calls, they are also supported, although
+  they require explicit parentheses around:
+
+      iex> fun = &(&1 + &2)
+      iex> fun.(1, 2)
+      3
+
+  And even more complex call expressions:
+
+      iex> fun = &(&1 + &2 + &3)
+      iex> fun.(1, 2, 3)
+      6
+
+  Remember tuple and lists are represented as calls in the AST and
+  therefore are also allowed:
+
+      iex> fun = &{&1, &2}
+      iex> fun.(1, 2)
+      { 1, 2 }
+
+      iex> fun = &[&1|&2]
+      iex> fun.(1, 2)
+      [1|2]
+
+  Anything that is not a call is not allowed though, examples:
+
+      # An atom is not a call
+      &:foo
+
+      # A var is not a call
+      var = 1
+      &var
+
+      # A block expression is not a call
+      &(foo(&1, &2); &3 + &4)
+
   """
   name = :&
   defmacro unquote(name)(expr)

lib/elixir/src/elixir_dispatch.erl

@@ -46,8 +46,12 @@ import_function(Meta, Name, Arity, S) ->
     { import, Receiver } ->
       require_function(Meta, Receiver, Name, Arity, S);
     nomatch ->
-      elixir_tracker:record_local(Tuple, S#elixir_scope.module, S#elixir_scope.function),
-      { { 'fun', ?line(Meta), { function, Name, Arity } }, S }
+      case elixir_import:special_form(Name, Arity) of
+        true  -> false;
+        false ->
+          elixir_tracker:record_local(Tuple, S#elixir_scope.module, S#elixir_scope.function),
+          { { 'fun', ?line(Meta), { function, Name, Arity } }, S }
+      end
   end.
 
 require_function(Meta, Receiver, Name, Arity, S) ->

lib/elixir/src/elixir_fn.erl

@@ -42,6 +42,101 @@ capture(Meta, { '/', _, [{ F, _, C }, A] }, S) when is_atom(F), is_integer(A), i
     Else  -> Else
   end;
 
+capture(Meta, { { '.', _, [Left, Right] }, RemoteMeta, Args } = Expr, S) when is_atom(Right), is_list(Args) ->
+  { Mod, SE } = 'Elixir.Macro':expand_all(Left, elixir_scope:to_ex_env({ ?line(Meta), S }), S),
+
+  case is_atom(Mod) andalso is_sequential(Args) andalso
+       elixir_dispatch:require_function(RemoteMeta, Mod, Right, length(Args), SE) of
+    false -> do_capture(Meta, Expr, S);
+    Else  -> Else
+  end;
+
+capture(Meta, { '__block__', _, _ } = Expr, S) ->
+  Message = "invalid args for &, block expressions are not allowed, got: ~ts",
+  syntax_error(Meta, S#elixir_scope.file, Message, ['Elixir.Macro':to_string(Expr)]);
+
+capture(Meta, { Atom, ImportMeta, Args } = Expr, S) when is_atom(Atom), is_list(Args) ->
+  case is_sequential(Args) andalso
+       elixir_dispatch:import_function(ImportMeta, Atom, length(Args), S) of
+    false -> do_capture(Meta, Expr, S);
+    Else  -> Else
+  end;
+
+capture(Meta, { Left, Right }, S) ->
+  capture(Meta, { '{}', Meta, [Left, Right] }, S);
+
+capture(Meta, List, S) when is_list(List) ->
+  capture(Meta, { '[]', Meta, List }, S);
+
+capture(Meta, Arg, S) when is_integer(Arg) ->
+  compile_error(Meta, S#elixir_scope.file, "unhandled &~B outside of a capture", [Arg]);
+
 capture(Meta, Arg, S) ->
-  syntax_error(Meta, S#elixir_scope.file,
-               "invalid args for &: ~ts", ['Elixir.Macro':to_string(Arg)]).
+  invalid_capture(Meta, Arg, S).
+
+%% Helpers
+
+do_capture(Meta, Expr, S) ->
+  case do_escape(Expr, S, []) of
+    { _, _, [] } ->
+      invalid_capture(Meta, Expr, S);
+    { TExpr, TS, TDict } ->
+      TVars = validate(Meta, TDict, 1, S),
+      fn(Meta, [{ TVars, Meta, TExpr }], TS)
+  end.
+
+invalid_capture(Meta, Arg, S) ->
+  Message = "invalid args for &, expected an expression in the format of &Mod.fun/arity, "
+            "&local/arity or a capture containing at least one argument as &1, got: ~ts",
+  syntax_error(Meta, S#elixir_scope.file, Message, ['Elixir.Macro':to_string(Arg)]).
+
+validate(Meta, [{ Pos, Var }|T], Pos, S) ->
+  [Var|validate(Meta, T, Pos + 1, S)];
+
+validate(Meta, [{ Pos, _ }|_], Expected, S) ->
+  compile_error(Meta, S#elixir_scope.file, "capture &~B cannot be defined without &~B", [Pos, Expected]);
+
+validate(_Meta, [], _Pos, _S) ->
+  [].
+
+do_escape({ '&', Meta, [Pos] }, S, Dict) when is_integer(Pos) ->
+  case orddict:find(Pos, Dict) of
+    { ok, Var } ->
+      { Var, S, Dict };
+    error ->
+      { Var, SC } = elixir_scope:build_ex_var(?line(Meta), S),
+      { Var, SC, orddict:store(Pos, Var, Dict) }
+  end;
+
+do_escape({ '&', Meta, _ } = Arg, S, _Dict) ->
+  Message = "nested captures via & are not allowed: ~ts",
+  compile_error(Meta, S#elixir_scope.file, Message, ['Elixir.Macro':to_string(Arg)]);
+
+do_escape({ Left, Meta, Right }, S0, Dict0) ->
+  { TLeft, S1, Dict1 }  = do_escape(Left, S0, Dict0),
+  { TRight, S2, Dict2 } = do_escape(Right, S1, Dict1),
+  { { TLeft, Meta, TRight }, S2, Dict2 };
+
+do_escape({ Left, Right }, S0, Dict0) ->
+  { TLeft, S1, Dict1 }  = do_escape(Left, S0, Dict0),
+  { TRight, S2, Dict2 } = do_escape(Right, S1, Dict1),
+  { { TLeft, TRight }, S2, Dict2 };
+
+do_escape(List, S, Dict) when is_list(List) ->
+  do_escape_list(List, S, Dict, []);
+
+do_escape(Other, S, Dict) ->
+  { Other, S, Dict }.
+
+do_escape_list([H|T], S, Dict, Acc) ->
+  { TH, TS, TDict } = do_escape(H, S, Dict),
+  do_escape_list(T, TS, TDict, [TH|Acc]);
+
+do_escape_list([], S, Dict, Acc) ->
+  { lists:reverse(Acc), S, Dict }.
+
+is_sequential(List) -> is_sequential(List, 1).
+is_sequential([{ '&', _, [Int] }|T], Int) ->
+  is_sequential(T, Int + 1);
+is_sequential([], Int) when Int > 1 -> true;
+is_sequential(_, _Int) -> false.

lib/elixir/src/elixir_import.erl

@@ -2,7 +2,7 @@
 %% in between local functions and imports.
 %% For imports dispatch, please check elixir_dispatch.
 -module(elixir_import).
--export([import/5, format_error/1]).
+-export([import/5, special_form/2, format_error/1]).
 -include("elixir.hrl").
 
 %% IMPORT HELPERS
@@ -139,11 +139,7 @@ get_optional_macros(Module)  ->
 %% VALIDATION HELPERS
 
 ensure_no_special_form_conflict(Meta, File, Key, [{Name,Arity}|T]) ->
-  Values = lists:filter(fun({X,Y}) ->
-    (Name == X) andalso ((Y == '*') orelse (Y == Arity))
-  end, special_form()),
-
-  case Values /= [] of
+  case special_form(Name, Arity) of
     true  ->
       Tuple = { special_form_conflict, { Key, Name, Arity } },
       elixir_errors:form_error(Meta, File, ?MODULE, Tuple);
@@ -198,44 +194,42 @@ remove_internals(Set) ->
   ordsets:del_element({ module_info, 1 },
     ordsets:del_element({ module_info, 0 }, Set)).
 
-%% Macros implemented in Erlang that are not importable.
-
-special_form() ->
-  [
-    {'&',1},
-    {'^',1},
-    {'=',2},
-    {'__op__',2},
-    {'__op__',3},
-    {'__scope__',2},
-    {'__block__','*'},
-    {'->','*'},
-    {'<<>>','*'},
-    {'{}','*'},
-    {'[]','*'},
-    {'alias',1},
-    {'alias',2},
-    {'require',1},
-    {'require',2},
-    {'import',1},
-    {'import',2},
-    {'import',3},
-    {'__ENV__',0},
-    {'__CALLER__',0},
-    {'__MODULE__',0},
-    {'__FILE__',0},
-    {'__DIR__',0},
-    {'__aliases__','*'},
-    {'quote',1},
-    {'quote',2},
-    {'unquote',1},
-    {'unquote_splicing',1},
-    {'fn','*'},
-    {'super','*'},
-    {'super?',0},
-    {'bc','*'},
-    {'lc','*'},
-    {'var!',1},
-    {'var!',2},
-    {'alias!',1}
-  ].
+%% Special forms
+
+special_form('&',1) -> true;
+special_form('^',1) -> true;
+special_form('=',2) -> true;
+special_form('__op__',2) -> true;
+special_form('__op__',3) -> true;
+special_form('__scope__',2) -> true;
+special_form('__block__',_) -> true;
+special_form('->',_) -> true;
+special_form('<<>>',_) -> true;
+special_form('{}',_) -> true;
+special_form('[]',_) -> true;
+special_form('alias',1) -> true;
+special_form('alias',2) -> true;
+special_form('require',1) -> true;
+special_form('require',2) -> true;
+special_form('import',1) -> true;
+special_form('import',2) -> true;
+special_form('import',3) -> true;
+special_form('__ENV__',0) -> true;
+special_form('__CALLER__',0) -> true;
+special_form('__MODULE__',0) -> true;
+special_form('__FILE__',0) -> true;
+special_form('__DIR__',0) -> true;
+special_form('__aliases__',_) -> true;
+special_form('quote',1) -> true;
+special_form('quote',2) -> true;
+special_form('unquote',1) -> true;
+special_form('unquote_splicing',1) -> true;
+special_form('fn',_) -> true;
+special_form('super',_) -> true;
+special_form('super?',0) -> true;
+special_form('bc',_) -> true;
+special_form('lc',_) -> true;
+special_form('var!',1) -> true;
+special_form('var!',2) -> true;
+special_form('alias!',1) -> true;
+special_form(_, _) -> false.

lib/elixir/test/elixir/kernel/fn_test.exs

@@ -2,26 +2,98 @@ Code.require_file "../test_helper.exs", __DIR__
 
 defmodule Kernel.FnTest do
   use ExUnit.Case, async: true
+  import CompileAssertion
 
   test "capture remote" do
     assert (&:erlang.atom_to_list/1).(:a) == 'a'
     assert (&Kernel.atom_to_list/1).(:a) == 'a'
 
     assert (&List.flatten/1).([[0]]) == [0]
     assert (&(List.flatten/1)).([[0]]) == [0]
+    assert (&List.flatten(&1)).([[0]]) == [0]
+    assert &List.flatten(&1) == &List.flatten/1
   end
 
   test "capture local" do
     assert (&atl/1).(:a) == 'a'
     assert (&(atl/1)).(:a) == 'a'
+    assert (&atl(&1)).(:a) == 'a'
   end
 
   test "capture imported" do
     assert (&atom_to_list/1).(:a) == 'a'
     assert (&(atom_to_list/1)).(:a) == 'a'
+    assert (&atom_to_list(&1)).(:a) == 'a'
+    assert &atom_to_list(&1) == &atom_to_list/1
+  end
+
+  test "local partial application" do
+    assert (&atb(&1, :utf8)).(:a) == "a"
+    assert (&atb(list_to_atom(&1), :utf8)).('a') == "a"
+  end
+
+  test "imported partial application" do
+    assert (&atom_to_binary(&1, :utf8)).(:a) == "a"
+    assert (&atom_to_binary(list_to_atom(&1), :utf8)).('a') == "a"
+  end
+
+  test "remote partial application" do
+    assert (&:erlang.atom_to_binary(&1, :utf8)).(:a) == "a"
+    assert (&:erlang.atom_to_binary(list_to_atom(&1), :utf8)).('a') == "a"
+  end
+
+  test "capture and partially apply tuples" do
+    assert (&{ &1, &2 }).(1, 2) == { 1, 2 }
+    assert (&{ &1, &2, &3 }).(1, 2, 3) == { 1, 2, 3 }
+
+    assert (&{ 1, &1 }).(2) == { 1, 2 }
+    assert (&{ 1, &1, &2 }).(2, 3) == { 1, 2, 3 }
+  end
+
+  test "capture and partially apply lists" do
+    assert (&[ &1, &2 ]).(1, 2) == [ 1, 2 ]
+    assert (&[ &1, &2, &3 ]).(1, 2, 3) == [ 1, 2, 3 ]
+
+    assert (&[ 1, &1 ]).(2) == [ 1, 2 ]
+    assert (&[ 1, &1, &2 ]).(2, 3) == [ 1, 2, 3 ]
+
+    assert (&[&1|&2]).(1, 2) == [1|2]
+  end
+
+  test "failure on non-continuous" do
+    assert_compile_fail CompileError, "nofile:1: capture &2 cannot be defined without &1", "&(&2)"
+  end
+
+  test "failure on integers" do
+    assert_compile_fail CompileError, "nofile:1: unhandled &1 outside of a capture", "&1"
+  end
+
+  test "failure on block" do
+    assert_compile_fail SyntaxError,
+      "nofile:1: invalid args for &, block expressions " <>
+      "are not allowed, got: (\n  1\n  2\n)",
+      "&(1;2)"
+  end
+
+  test "failure on other types" do
+    assert_compile_fail SyntaxError,
+      "nofile:1: invalid args for &, expected an expression in the format of &Mod.fun/arity, " <>
+      "&local/arity or a capture containing at least one argument as &1, got: :foo",
+      "&:foo"
+  end
+
+  test "failure when no captures" do
+    assert_compile_fail SyntaxError,
+      "nofile:1: invalid args for &, expected an expression in the format of &Mod.fun/arity, " <>
+      "&local/arity or a capture containing at least one argument as &1, got: foo()",
+      "&foo()"
   end
 
   defp atl(arg) do
     :erlang.atom_to_list arg
   end
-end
+
+  defp atb(arg, encoding) do
+    :erlang.atom_to_binary(arg, encoding)
+  end
+end

lib/elixir/test/elixir/kernel/quote_test.exs

@@ -298,7 +298,7 @@ defmodule Kernel.QuoteTest.ImportsHygieneTest do
 
   defmacrop get_bin_size_with_partial do
     quote do
-      size(&1).("hello")
+      (&size(&1)).("hello")
     end
   end