%   Module : lens
%   Author : Richard A. O'Keefe
%   Updated: 3-Dec-2015
%   Purpose: Implement Haskell-style Lenses for Erlang
%   SeeAlso: https://github.com/jlouis/erl-lenses

-module(lens).
-export([
    get/2, put/3, update/3, c/2, c/3, c/4, c/5, c/6, complete/2,
    id/0,
    hd/0, tl/0, index/1, where/1, all/1,
    gb_lookup/1, % gb_tree/1,
    integer/0,
    tuple/1,
    pair_fst/0, pair_snd/0,
    triple_fst/0, triple_snd/0, triple_thd/0
]).
-export([ % for lens.hrl
    nth/2, set_nth/3, update_nth/3,
    find/2, set_find/3, update_find/3,
    map2/3
]).
-export([ % for testing
    tc/0, td/0, df/2
]).

-compile({inline, [
    get/2, put/3, update/3, c/2, c/3, c/4, c/5, c/6,
    hd/0, tl/0, index/1, where/1,
    tuple/1,
    pair_fst/0, pair_snd/0,
    triple_fst/0, triple_snd/0, triple_thd/0
]}).

%   A lens is a triple of functions
%   { Get :: (A) -> B,  Put :: (A, B) -> A, Upd :: (A, (B) -> B) -> A. }.
%   They *should* satisfy the following laws:
%       Get(Put(A, B)) = B
%       Put(A, Get(A)) = A
%       Upd(A, Fun)    = Put(A, Fun(Get(A)))
%       Put(A, B)      = Upd(A, fun (_) -> B end)
%   and a meta-law:
%       No side effects except raising an exception for non-existent data.
%   The code below doesn't actually depend on these laws,
%   but if we expect lenses to act like field access and indexing
%   the laws had better be followed.

%   We can apply these operations in various ways to extract and
%   replace information, but we can also *compose* them.
%   See the df/2 example below, which composes 3 field selectors
%   to make something that can be got or put in what looks like one step.

get({G,_,_}, X) -> G(X).

put({_,P,_}, X, Y) -> P(X, Y).

update({_,_,U}, X, F) -> U(X, F).

complete(G, P) ->
    { G
    , P
    , fun (X, F) -> P(X, F(G(X))) end
    }.

%   id() is the identity lens.
%   I owe this to Jesper Louis Anderson, who wrote and published
%   a lens implementation for Erlang back in July 2012.
%   We expect that c(id(), L) = c(L, id()) = L.

id() ->
    { fun (X)    -> X end
    , fun (_, Y) -> Y end
    , fun (X, F) -> F(X) end
    }.

%   get(hd(), X) = head(X)

hd() ->
    { fun ([H|_])    -> H end
    , fun ([_|T], H) -> [H|T] end
    , fun ([H|T], F) -> [F(H)|T] end
    }.

%   get(tl(), X) = tail(X)

tl() ->
    { fun ([_|T])    -> T end
    , fun ([H|_], T) -> [H|T] end
    , fun ([H|T], F) -> [H|F(T)] end
    }.

nth(1, [H|_]) -> H;
nth(N, [_|T]) when N > 1 ->
    nth(N-1, T).

set_nth(1, [_|T], X) -> [X|T];
set_nth(N, [H|T], X) when N > 1 ->
    [H|set_nth(N-1, T, X)].

update_nth(1, [X|T], F) -> [F(X)|T];
update_nth(N, [H|T], F) when N > 1 ->
    [H|update_nth(N-1, T, F)].

%   get(index(N), L) = the Nth element of  L, 1-origin.
%   It is an error if there is no such element.

index(N) ->
    { fun (Xs)    -> nth(N, Xs) end
    , fun (Xs, Y) -> set_nth(N, Xs, Y) end
    , fun (Xs, F) -> update_nth(N, Xs, F) end
    }.

find(P, [H|T]) ->
    case P(H)
      of true  -> H
       ; false -> find(P, T)
    end.

set_find(P, [H|T], X) ->
    case P(H)
      of true  -> [X|T]
       ; false -> [H|set_find(P, T, X)]
    end.

update_find(P, [H|T], F) ->
    case P(H)
      of true  -> [F(H)|T]
       ; false -> [H|update_find(P, T, F)]
    end.       

%   get(where(P), L) is the first element X of L for which P(X).
%   It is an error if there is no such element.

where(P) ->
    { fun (Xs)    -> find(P, Xs) end
    , fun (Xs, Y) -> set_find(P, Xs, Y) end
    , fun (Xs, F) -> update_find(P, Xs, F) end
    }.

%   get(tuple(N), T) = element(N, T).
%   Its real intended use is tuple(#<record>.<field>).

tuple(N) ->
    { fun (T)    -> element(N, T) end
    , fun (T, Y) -> setelement(N, T, Y) end
    , fun (T, F) -> setelement(N, T, F(element(N, T))) end
    }.


pair_fst() ->
    { fun ({A,_})    -> A end
    , fun ({_,B}, A) -> {A,B} end
    , fun ({A,B}, F) -> {F(A),B} end
    }.

pair_snd() ->
    { fun ({_,B})    -> B end
    , fun ({A,_}, B) -> {A,B} end
    , fun ({A,B}, F) -> {A,F(B)} end
    }.

triple_fst() ->
    { fun ({A,_,_})    -> A end
    , fun ({_,B,C}, A) -> {A,B,C} end
    , fun ({A,B,C}, F) -> {F(A),B,C} end
    }.

triple_snd() ->
    { fun ({_,B,_})    -> B end
    , fun ({A,_,C}, B) -> {A,B,C} end
    , fun ({A,B,C}, F) -> {A,F(B),C} end
    }.

triple_thd() ->
    { fun ({_,_,C})    -> C end
    , fun ({A,B,_}, C) -> {A,B,C} end
    , fun ({A,B,C}, F) -> {A,B,F(C)} end
    }.

%   get(integer(), N) is N as a list of character codes.
%   This is a jeux d'esprit to point out that the information
%   we extract from a value does not have to be stored in it.

integer() ->
    { fun (N) -> integer_to_list(N) end
    , fun (_, L) -> list_to_integer(L) end
    , fun (N, F) -> list_to_integer(F(integer_to_list(N))) end
    }.

%   We should also provide a lens for gb_trees.
%   This code is commented out because the updatef function is missing.
%
%gb_tree(Key) ->
%   { fun (Tree)    -> gb_trees:get(Key, Tree) end
%   , fun (Tree, B) -> gb_trees:update(Key, Tree, B) end
%   , fun (Tree, F) -> gb_trees:updatef(Key, Tree, F) end
%   }/

%  gb_lookup(K) -> lens(gb_tree(K,V), none | {value,V}).

gb_lookup(Key) ->
    { fun (Tree) -> gb_trees:lookup(Key, Tree) end
    , fun (Tree, none)      -> gb_trees:delete_any(Key, Tree)
        ; (Tree, {value,V}) -> gb_trees:enter(Key, V, Tree)
      end
    , fun (Tree, F) ->
          case gb_trees:lookup(Key, Tree)
            of none -> Tree
             ; {value,V} -> gb_trees:update(Key, F(V), Tree)
          end
      end
    }.


%   c(Lens1, Lens2[, Lens3[, Lens4[, Lens5[, Lens6]]]])
%   compose 2 to 6 lens.  Composition is associative with
%   identity id, but not commutative.  These definitions
%   were generated by a program.

c({G1,P1,U1}, {G2,P2,U2}) ->
    { fun (F0) -> G2(G1(F0)) end
    , fun (F0, R2) ->
        F1 = G1(F0),
        R1 = P2(F1, R2),
        P1(F0, R1)
      end
    , fun (F0, UF) ->
        U1(F0, fun (F1) -> 
          U2(F1, UF)
        end)
      end
    }.

c({G1,P1,U1}, {G2,P2,U2}, {G3,P3,U3}) ->
    { fun (F0) -> G3(G2(G1(F0))) end
    , fun (F0, R3) ->
        F1 = G1(F0),
        F2 = G2(F1),
        R2 = P3(F2, R3),
        R1 = P2(F1, R2),
        P1(F0, R1)
      end
    , fun (F0, UF) ->
        U1(F0, fun (F1) -> 
          U2(F1, fun (F2) -> 
            U3(F2, UF)
          end)
        end)
      end
    }.

c({G1,P1,U1}, {G2,P2,U2}, {G3,P3,U3}, {G4,P4,U4}) ->
    { fun (F0) -> G4(G3(G2(G1(F0)))) end
    , fun (F0, R4) ->
        F1 = G1(F0),
        F2 = G2(F1),
        F3 = G3(F2),
        R3 = P4(F3, R4),
        R2 = P3(F2, R3),
        R1 = P2(F1, R2),
        P1(F0, R1)
      end
    , fun (F0, UF) ->
        U1(F0, fun (F1) -> 
          U2(F1, fun (F2) -> 
            U3(F2, fun (F3) -> 
              U4(F3, UF)
            end)
          end)
        end)
      end
    }.

c({G1,P1,U1}, {G2,P2,U2}, {G3,P3,U3}, {G4,P4,U4}, {G5,P5,U5}) ->
    { fun (F0) -> G5(G4(G3(G2(G1(F0))))) end
    , fun (F0, R5) ->
        F1 = G1(F0),
        F2 = G2(F1),
        F3 = G3(F2),
        F4 = G4(F3),
        R4 = P5(F4, R5),
        R3 = P4(F3, R4),
        R2 = P3(F2, R3),
        R1 = P2(F1, R2),
        P1(F0, R1)
      end
    , fun (F0, UF) ->
        U1(F0, fun (F1) -> 
          U2(F1, fun (F2) -> 
            U3(F2, fun (F3) -> 
              U4(F3, fun (F4) -> 
                U5(F4, UF)
              end)
            end)
          end)
        end)
      end
    }.

c({G1,P1,U1}, {G2,P2,U2}, {G3,P3,U3}, {G4,P4,U4}, {G5,P5,U5}, {G6,P6,U6}) ->
    { fun (F0) -> G6(G5(G4(G3(G2(G1(F0)))))) end
    , fun (F0, R6) ->
        F1 = G1(F0),
        F2 = G2(F1),
        F3 = G3(F2),
        F4 = G4(F3),
        F5 = G5(F4),
        R5 = P6(F5, R6),
        R4 = P5(F4, R5),
        R3 = P4(F3, R4),
        R2 = P3(F2, R3),
        R1 = P2(F1, R2),
        P1(F0, R1)
      end
    , fun (F0, UF) ->
        U1(F0, fun (F1) -> 
          U2(F1, fun (F2) -> 
            U3(F2, fun (F3) -> 
              U4(F3, fun (F4) -> 
                U5(F4, fun (F5) -> 
                  U6(F5, UF)
                end)
              end)
            end)
          end)
        end)
      end
    }.

map2(F, [X|Xs], [Y|Ys]) ->
    [F(X, Y) | map2(F, Xs, Ys)];
map2(_, [], []) ->
    [].

all({G,P,U}) ->
    { fun (Xs) -> [G(X) || X <- Xs] end
    , fun (Xs, Ys) -> map2(P, Xs, Ys) end
    , fun (Xs, F) -> [U(X, F) || X <- Xs] end
    }.

%-----------------------------------------------------------------------
% Record fields #rec.field are integers but lenses would be nicer.
%-----------------------------------------------------------------------

-record(a, {p,q,r,s}).

df(R, Y) ->
  % R.p.q.r <- Y
    put(c(tuple(#a.p), tuple(#a.q), tuple(#a.r)), R, Y).

%-----------------------------------------------------------------------
% Examples.
%-----------------------------------------------------------------------

tc() ->
    c(pair_fst(), triple_thd()).

td() ->
    {{1,2,3},4}.

% 1> c(lens).
% {ok,lens}
% 2> lens:df({a,{a,p2,{a,p3,q3,r3,s3},r2,s2},q1,r1,s1}, 42).
% {a,{a,p2,{a,p3,q3,42,s3},r2,s2},q1,r1,s1}
% 3> lens:get(lens:tc(), lens:td()).
% 3
% 4> lens:put(lens:tc(), lens:td(), 5).
% {{1,2,5},4}
% 5> lens:update(lens:tc(), lens:td(), fun (N) -> 10*N-1 end).
% {{1,2,29},4}

%-----------------------------------------------------------------------
% Why we want cross-module inlining.
%-----------------------------------------------------------------------
%   consider
%   lens:update(lens:triple_snd(), Triple, fun (X) -> X+1 end).
%   (1) Expand lens:triple_snd/0 inline
%   lens:update({G,P,U}, Triple, fun (X) -> X + 1 end)
%      where G = fun ({_,B,_}) -> B end
%        and P = fun ({A,_,C}, B) -> {A,B,C} end
%        and U = fun ({A,B,C}, F) -> {A,F(B),C} end
%   (2) Expand lens:update/3 inline
%   U(Triple, f(X) -> X + 1 end
%      where U = fun ({A,B,C}, F) -> {A,F(B),C} end
%   (3) Expand U(X) inline
%   {A,B,C} = Triple, F = fun (X) -> X + 1 end, {A,F(B),C}
%   (4) expand F(X) inline
%   {A,B,C} = Triple, {A,B+1,C}
%   It's hard to improve this.

%-----------------------------------------------------------------------
% Lists.
%-----------------------------------------------------------------------
%
% We can index into a matrix represented as a list of lists:
% index2(Row, Col) -> c(list_lens(Row), list_lens(Col)).
% a(i,j) := x => put(index2(I,J), A, X).
% Easy to do, but the high cost of indexing into lists means
% this will never be a good idea.