Rewrite records accessors when possible

Author: José Valim

lib/elixir/lib/kernel/record_rewriter.ex

@@ -38,13 +38,75 @@ defmodule Kernel.RecordRewriter do
     optimize_body(t, new_dict, [new_expr|acc])
   end
 
+  ## Record helpers
+
+  defp record_fields(record) do
+    if Code.ensure_loaded?(record) && function_exported?(record, :__record__, 1) do
+      try do
+        fields      = lc { k, _ } inlist record.__record__(:fields), do: k
+        optimizable = record.__record__(:optimizable)
+        { fields, optimizable }
+      rescue
+        [UndefinedFunctionError, FunctionClauseError] -> { [], [] }
+      end
+    end
+  end
+
+  defp record_field_info(function) do
+    case atom_to_list(function) do
+      'update_' ++ field -> { :update, list_to_atom(function) }
+      _ -> { :accessor, function }
+    end
+  end
+
+  defp optimize_call(line, { record, _ } = res, left, { :atom, _, function }, args) do
+    { fields, optimizable } = record_fields(record)
+
+    if List.member?(optimizable, { function, length(args) + 1 }) do
+      { kind, field } = record_field_info(function)
+      if index = Enum.find_index(fields, field == &1) do
+        optimize_call(line, res, kind, index, left, args)
+      end
+    end
+  end
+
+  defp optimize_call(_line, _res, _left, _right, _args) do
+    nil
+  end
+
+  defp optimize_call(line, _res, :accessor, index, left, []) do
+    call = { :call, line,
+      { :remote, line, { :atom, 0, :erlang }, { :atom, 0, :element } },
+      [{ :integer, 0, index + 2 }, left]
+    }
+    { call, nil }
+  end
+
+  defp optimize_call(line, res, :accessor, index, left, [arg]) do
+    call = { :call, line,
+      { :remote, line, { :atom, 0, :erlang }, { :atom, 0, :setelement } },
+      [{ :integer, 0, index + 2 }, left, arg]
+    }
+    { call, res }
+  end
+
+  defp optimize_call(_line, _res, :update, _index, _left, [_arg]) do
+    nil
+  end
+
   ## Expr
 
   defp optimize_expr({ :call, call_line, { :remote, line, left, right }, args }, dict) do
-    { left, dict, _ }  = optimize_expr(left, dict)
+    { left, dict, res } = optimize_expr(left, dict)
     { right, dict, _ } = optimize_expr(right, dict)
     { args, dict } = optimize_args(args, dict)
-    { { :call, call_line, { :remote, line, left, right }, args }, dict, nil }
+
+    case optimize_call(call_line, res, left, right, args) do
+      { call, call_res } ->
+        { call, dict, call_res }
+      nil ->
+        { { :call, call_line, { :remote, line, left, right }, args }, dict, nil }
+    end
   end
 
   defp optimize_expr({ :call, line, expr, args }, dict) do
@@ -212,21 +274,19 @@ defmodule Kernel.RecordRewriter do
   end
 
   defp assign_vars([key|t], dict, { value, _ } = res) when is_atom(key) and value != nil do
-    if is_record?(value) do
-      dict =
-        case :orddict.find(key, dict) do
-          { :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, res, dict)
-        end
-    end
+    dict =
+      case :orddict.find(key, dict) do
+        { :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, res, dict)
+      end
 
     assign_vars t, dict, res
   end
@@ -305,11 +365,4 @@ defmodule Kernel.RecordRewriter do
   defp join_result([], res) do
     res
   end
-
-  ## Record helpers
-
-  # TODO: Implement proper record check
-  defp is_record?(_h) do
-    true
-  end
 end

lib/elixir/lib/macro/env.ex

@@ -16,7 +16,7 @@ defmodule Macro.Env do
     functions: functions, macros: macros]
 
   Record.deffunctions(fields, __MODULE__)
-  Record.deftypes(fields, types, __ENV__)
+  Record.deftypes(fields, types, __MODULE__)
 
   @moduledoc """
   A record that contains compile time environment information,

lib/elixir/lib/module.ex

@@ -1,6 +1,12 @@
 defmodule Module do
   require :ets, as: ETS
 
+  defmacrop is_env(env) do
+    quote do
+      is_tuple(unquote(env)) and size(unquote(env)) > 1 and elem(unquote(env), 0) == Macro.Env
+    end
+  end
+
   @moduledoc """
   This module provides many functions to deal with modules during
   compilation time. It allows a developer to dynamically attach
@@ -51,11 +57,11 @@ defmodule Module do
   """
   def eval_quoted(module, quoted, binding // [], opts // [])
 
-  def eval_quoted(Macro.Env[module: module] = env, quoted, binding, opts) do
-    eval_quoted(module, quoted, binding, Keyword.merge(env.to_keywords, opts))
+  def eval_quoted(env, quoted, binding, opts) when is_env(env) do
+    eval_quoted(env.module, quoted, binding, Keyword.merge(env.to_keywords, opts))
   end
 
-  def eval_quoted(module, quoted, binding, Macro.Env[] = env) do
+  def eval_quoted(module, quoted, binding, env) when is_env(env) do
     eval_quoted(module, quoted, binding, env.to_keywords)
   end
 
@@ -95,12 +101,12 @@ defmodule Module do
   """
   def create(module, quoted, opts // [])
 
-  def create(module, quoted, Macro.Env[] = env) do
+  def create(module, quoted, env) when is_env(env) do
     create(module, quoted, env.to_keywords)
   end
 
   def create(module, quoted, opts) when is_atom(module) do
-    line = opts[:line] || 1
+    line = Keyword.get(opts, :line, 1)
     :elixir_module.compile(line, module, quoted, [], :elixir.scope_for_eval(opts))
   end
 
@@ -578,7 +584,7 @@ defmodule Module do
     atom
   end
 
-  defp normalize_attribute(:file, Macro.Env[file: file, line: line]),       do: { binary_to_list(file), line}
+  defp normalize_attribute(:file, env) when is_env(env),      do: { binary_to_list(env.file), env.line}
   defp normalize_attribute(:file, { binary, line }) when is_binary(binary), do: { binary_to_list(binary), line }
   defp normalize_attribute(:file, other) when not is_tuple(other),          do: normalize_attribute(:file, { other, 1 })
 

lib/elixir/lib/record.ex

@@ -73,17 +73,18 @@ defmodule Record do
       reflection(escaped),
       initializer(escaped),
       indexes(escaped),
-      accessors(values, 1, []),
       conversions(values),
       updater(values),
-      extensions(values, 1, [], Record.Extensions)
+      extensions(values, 1, [], Record.Extensions),
+      accessors(values),
     ]
 
+    contents = [quote(do: @__record__ unquote(escaped))|contents]
+
     # Special case for bootstraping purposes
     if env == Macro.Env do
-      :elixir_module.eval_quoted(env, contents, [], [])
+      Module.eval_quoted(env, contents, [], [])
     else
-      contents = [quote(do: @__record__ unquote(escaped))|contents]
       Module.eval_quoted(env.module, contents, [], env.location)
     end
   end
@@ -103,8 +104,9 @@ defmodule Record do
       accessor_specs(values, 1, [])
     ]
 
-    # Special case for bootstraping purposes
-    if :erlang.function_exported(Module, :eval_quoted, 2) do
+    if env == Macro.Env do
+      Module.eval_quoted(env, contents, [], [])
+    else
       Module.eval_quoted(env.module, contents, [], env.location)
     end
   end
@@ -152,6 +154,34 @@ defmodule Record do
     Module.eval_quoted(env.module, contents, [], env.location)
   end
 
+  ## Callbacks
+
+  # Store all optimizable fields in the record as well
+  @doc false
+  defmacro __before_compile__(_) do
+    quote do
+      def __record__(:optimizable), do: @record_optimizable
+    end
+  end
+
+  # Store fields that can be optimized and that cannot be
+  # optimized as they are overriden
+  @doc false
+  def __on_definition__(env, kind, name, args, _guards, _body) do
+    tuple     = { name, length(args) }
+    module    = env.module
+    functions = Module.get_attribute(module, :record_optimizable)
+
+    functions =
+      if kind in [:def] and Module.get_attribute(module, :record_optimized) do
+        [tuple|functions]
+      else
+        List.delete(functions, tuple)
+      end
+
+    Module.put_attribute(module, :record_optimizable, functions)
+  end
+
   # Implements the access macro used by records.
   # It returns a quoted expression that defines
   # a record or a match in case the record is
@@ -414,11 +444,20 @@ defmodule Record do
   #       setelem(record, 2, callback.(elem(record, 2)))
   #     end
   #
-  defp accessors([{ :__exception__, _ }|t], 1, acc) do
-    accessors(t, 2, acc)
+  defp accessors(values) do
+    [ quote do
+        @record_optimized true
+        @record_optimizable []
+        @before_compile { unquote(__MODULE__), :__before_compile__ }
+        @on_definition { unquote(__MODULE__), :__on_definition__ }
+      end | accessors(values, 1) ]
+  end
+
+  defp accessors([{ :__exception__, _ }|t], 1) do
+    accessors(t, 2)
   end
 
-  defp accessors([{ key, _default }|t], i, acc) do
+  defp accessors([{ key, _default }|t], i) do
     update = binary_to_atom "update_" <> atom_to_binary(key)
 
     contents = quote do
@@ -436,10 +475,12 @@ defmodule Record do
       end
     end
 
-    accessors(t, i + 1, [contents | acc])
+    [contents|accessors(t, i + 1)]
   end
 
-  defp accessors([], _i, acc), do: acc
+  defp accessors([], _i) do
+    [quote do: @record_optimized false]
+  end
 
   # Define an updater method that receives a
   # keyword list and updates the record.

lib/elixir/src/elixir_compiler.erl

@@ -159,11 +159,11 @@ core_main() ->
     "lib/elixir/lib/keyword.ex",
     "lib/elixir/lib/list.ex",
     "lib/elixir/lib/kernel/typespec.ex",    
+    "lib/elixir/lib/module.ex",
     "lib/elixir/lib/record.ex",
     "lib/elixir/lib/record/extractor.ex",
     "lib/elixir/lib/macro.ex",
     "lib/elixir/lib/macro/env.ex",
-    "lib/elixir/lib/module.ex",
     "lib/elixir/lib/code.ex",
     "lib/elixir/lib/protocol.ex",
     "lib/elixir/lib/enum.ex",

lib/elixir/src/elixir_translator.erl

@@ -500,11 +500,9 @@ translate_apply(Line, TLeft, TRight, Args, S, SL, SR) ->
   Optimize = case (Args == []) orelse lists:last(Args) of
     { '|', _, _ } -> false;
     _ ->
-      case { TLeft, TRight } of
-        { { Kind, _, _ }, { atom, _, _ } } when Kind == var; Kind == tuple; Kind == atom ->
-          true;
-        _ ->
-          false
+      case TRight of
+        { atom, _, _ } -> true;
+        _ -> false
       end
   end,
 

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

@@ -1,5 +1,13 @@
 Code.require_file "../../test_helper.exs", __FILE__
 
+defrecord BadRange, first: 0, last: 0 do
+  defoverridable [first: 1]
+
+  def first(_) do
+    :not_optimizable
+  end
+end
+
 defmodule Kernel.RecordRewriterTest do
   use ExUnit.Case, async: true
 
@@ -21,7 +29,7 @@ defmodule Kernel.RecordRewriterTest do
     { clause, dict, res }
   end
 
-  ## Dictionary tests
+  ## Inference tests
 
   test "simple atom" do
     clause = clause(fn -> :foo end)
@@ -198,4 +206,37 @@ defmodule Kernel.RecordRewriterTest do
     clause = clause(fn -> try do x = Macro.Env[]; x; after x = Macro.Env[]; end end)
     assert optimize_clause(clause) == { clause, [], { Macro.Env, nil } }
   end
+
+  ## Rewrite tests
+
+  test "getter call is rewriten" do
+    { clause, rewriten } =
+      { clause(fn(x = Range[]) -> x.first end), clause(fn(x = Range[]) -> :erlang.element(2, x) end) }
+
+    assert optimize_clause(clause) == { rewriten, [x: Range], nil }
+  end
+
+  test "setter call is rewriten" do
+    { clause, rewriten } =
+      { clause(fn(x = Range[]) -> x.first(:first) end), clause(fn(x = Range[]) -> :erlang.setelement(2, x, :first) end) }
+
+    assert optimize_clause(clause) == { rewriten, [x: Range], { Range, nil } }
+  end
+
+  test "nested setter call is rewriten" do
+    { clause, rewriten } =
+      { clause(fn(x = Range[]) -> x.first(:first).last(:last) end), clause(fn(x = Range[]) -> :erlang.setelement(3, :erlang.setelement(2, x, :first), :last) end) }
+
+    assert optimize_clause(clause) == { rewriten, [x: Range], { Range, nil } }
+  end
+
+  test "noop for unknown fields" do
+    clause = clause(fn(x = Range[]) -> x.unknown end)
+    assert optimize_clause(clause) == { clause, [x: Range], nil }
+  end
+
+  test "noop for rewriten fields" do
+    clause = clause(fn(x = BadRange[]) -> x.first end)
+    assert optimize_clause(clause) == { clause, [x: BadRange], nil }
+  end
 end