%% @private -module(gradualizer_highlight). -export([prettyprint_and_highlight/3, highlight_in_source/3]). -include("gradualizer.hrl"). -define(context_before, 2). -define(context_after, 0). -define(marker_char, $^). -define(color_text, "\e[36m"). % 36=cyan -define(color_marker, "\e[36m"). -define(color_end, "\e[0m"). %% 31 = Red, 91 = Bright red, 4 = Underline, 1 = Bold %% erl_parse:erl_parse_tree() is documented but not exported :-( -type erl_parse_tree() :: erl_parse:abstract_clause() | erl_parse:abstract_expr() | erl_parse:abstract_form() | erl_parse:abstract_type(). %% Pretty-prints and highlights a node in an AST. `AstNode' must be a node %% existing in the list `AstContext` (forms). %% %% To highlight a node in an AST without having the original source code: %% Pretty-print the AST and parse again. Then, find the corresponding node %% in the new AST. Find its location. Find its length by pretty-printing, %% tokenizing and then checking the location and length of the last token. %% Then print it all in a fancy way with the node highlighted. -spec prettyprint_and_highlight(AstNode :: erl_parse_tree(), AstContext :: [erl_parse_tree()], Color :: boolean()) -> iolist(). prettyprint_and_highlight(AstNode, AstContext, Color) -> TrimmedCtx = trim_context(AstNode, AstContext), {Pretty, NewNode, _NewCtx} = case recreate_source(AstNode, TrimmedCtx) of not_found when TrimmedCtx /= AstContext -> %% Trimming may have discarded the containing form. %% This can happen if line numbers in the AST are %% not in order. Try without the full context. recreate_source(AstNode, AstContext); {_,_,_} = Found -> Found end, highlight_in_source(NewNode, Pretty, Color). %% Recreates the source code of abstract forms and a node within it by %% pretty-printing it. The node is localized in the pretty-printed %% text by parsing the pretty-printed source code and matching them %% side by side to find the corresponding AstNode in the pretty-printed %% source code. recreate_source(AstNode, AstCtx) -> Pretty = prettyprint_forms(AstCtx), NewCtx = parse_into_list_of_forms(Pretty), case find_in_asts(AstNode, AstCtx, NewCtx) of {found, NewNode} -> {Pretty, NewNode, NewCtx}; not_found -> %% Could be that the line numbers in the original AST are not %% in ascending order. not_found end. %% Highlights a node in the AST in a piece of source code. The node must %% be annotated with line and column information for this to work. Returns %% an inlist(). highlight_in_source(AstNode, Source, Color) -> {Start, End} = find_start_and_end_location_in_source(AstNode, Source), highlight_text(Source, Start, End, Color). %% Finds the start and end location of an abstract node in source code. %% This is done by pretty-printing, tokenizing, counting the tokens and %% then finding the corresponding tokens in the tokenized source code. -spec find_start_and_end_location_in_source(Node :: erl_parse_tree(), Source :: string()) -> {Start :: erl_anno:location(), End :: erl_anno:location()}. find_start_and_end_location_in_source(Node, Source) -> StartLoc = min_location_in_ast(Node), %% Get the location after the last token of Node %% TODO: Generalize to not only Erlang (e.g. configurable fun) {ok, Tokens, _} = erl_scan:string(Source, {1,1}, [text]), TokensStartingAtNode = lists:dropwhile(fun (T) -> erl_anno:location(element(2,T)) < StartLoc end, Tokens), NumTokens = num_tokens(Node), [LastToken] = lists:sublist(TokensStartingAtNode, NumTokens, 1), EndLoc = {_,_} = erl_scan:end_location(LastToken), {StartLoc, EndLoc}. min_location_in_ast(AstNode) -> erl_parse:fold_anno(fun (Anno, Min) -> Loc = erl_anno:location(Anno), min(Min, Loc) end, erl_anno:location(element(2, AstNode)), AstNode). %% Returns the number of tokens in an AST node. %% This is done by pretty-printing it and tokensizing it again. -spec num_tokens(Ast :: erl_parse_tree()) -> pos_integer(). num_tokens(Ast) -> Pretty = erl_prettypr:format(Ast), {ok, Tokens, _End} = erl_scan:string(Pretty, {1,1}), ?assert_type(length(Tokens), pos_integer()). %% Find the form containing the searched node and sometimes a node before %% as extra context. trim_context(Node, Forms) -> case erl_anno:line(element(2, Node)) of 0 -> Forms; % generated code or something... Line -> Forms1 = lists:reverse(Forms), %% The containing form and the preceding forms reversed FormAndPreforms = lists:dropwhile(fun (Form) -> L = erl_anno:line(element(2, Form)), L > Line end, Forms1), case FormAndPreforms of [] -> %% Line numbers seem to be wrong Forms; [FoundForm|PreForms1] -> %% Include any preceding forms starting within 2 lines %% above our searched form PreForms2 = lists:takewhile(fun (Form) -> L = erl_anno:line(element(2, Form)), L >= Line - 2 end, PreForms1), lists:reverse([FoundForm|PreForms2]) end end. -spec prettyprint_forms(Forms :: [erl_parse_tree()]) -> string(). prettyprint_forms(Forms) -> lists:flatten([[erl_prettypr:format(Form), "\n"] || Form <- Forms]). -spec parse_into_list_of_forms(Forms :: string()) -> [erl_parse_tree()]. parse_into_list_of_forms(Pretty) -> case merl:quote({1,1}, Pretty) of Forms when is_list(Forms) -> Forms; Form -> [Form] end. %% Highlights the text between two locations (line and column) in a text. highlight_text(Pretty, StartLoc, EndLoc, Color) -> Lines = re:split(Pretty, "\\n", [{return, list}]), HiLines = lists:flatmap(fun ({Line, LineNo}) -> highlight_line(Line, LineNo, StartLoc, EndLoc, Color) end, lists:zip(Lines, lists:seq(1, length(Lines)))), %% trim leading empty lines HiLines1 = lists:dropwhile(fun ("") -> true; (_) -> false end, HiLines), [[Line, "\n"] || Line <- HiLines1]. %% Highlights a line between two locations by returning a list of lines: %% If the line is highlighted, one highliged line and a marker line is %% returned. If a line is too far from the highlighted area, no lines are %% returned. Otherwise, a list containing the unchanged line is returned. -spec highlight_line(Line :: string(), LineNo :: erl_anno:line(), Start :: erl_anno:location(), End :: erl_anno:location(), Color :: boolean()) -> [string()]. highlight_line(_Line, N, {L1, _}, {L2, _}, _Color) when N < L1 - ?context_before; N > L2 + ?context_after -> []; % Not included highlight_line(Line, N, {L1, _}, {L2, _}, _Color) when N < L1; N > L2 -> [Line]; % Not highlighted context highlight_line(Line, N, {N, C1}, {N, C2}, Color) -> %% Only this line highlighted, from C1 to C2 color_and_mark_line(Line, C1, C2, Color); highlight_line(Line, N, {L1, C1}, {L2, _C2}, Color) when L1 == N, N < L2 -> %% First line color_and_mark_line(Line, C1, length(Line) + 1, Color); highlight_line(Line, N, {L1, _C1}, {L2, C2}, Color) when L1 < N, N == L2 -> %% Last line color_and_mark_line(Line, step_spaces(Line, 1), C2, Color); highlight_line(Line, N, {L1, _}, {L2, __}, Color) when L1 < N, N < L2 -> %% Internal line color_and_mark_line(Line, step_spaces(Line, 1), length(Line) + 1, Color). -spec step_spaces(string(), erl_anno:column()) -> erl_anno:column(). step_spaces([$\t | Str], Col) -> step_spaces(Str, Col + 1); step_spaces([$\s | Str], Col) -> step_spaces(Str, Col + 1); step_spaces(_Str, Col) -> Col. blank([$\t | Str]) -> [$\t | blank(Str)]; blank([_Ch | Str]) -> [$\s | blank(Str)]; blank("") -> "". %% Color the string from column C1 to (not including) C2. Returns %% a list containing the colored line and a line with ^^^^ markers. -spec color_and_mark_line(Line, StartCol, EndCol, Color) -> [string()] when Line :: string(), StartCol :: erl_anno:column(), EndCol :: erl_anno:column(), Color :: boolean(). color_and_mark_line(Line, C1, C2, Color) when C2 >= C1 -> {Pre, Rest} = lists:split(C1 - 1, Line), HighlightLen = ?assert_type(C2 - C1, non_neg_integer()), {Mid, Post} = lists:split(HighlightLen, Rest), {ColorText, ColorMarker, ColorEnd} = case Color of true -> {?color_text, ?color_marker, ?color_end}; false -> {"", "", ""} end, [Pre ++ ColorText ++ Mid ++ ColorEnd ++ Post, blank(Pre) ++ ColorMarker ++ lists:duplicate(HighlightLen, ?marker_char) ++ ColorEnd]. %% Finds a node in an AST and return its corresponding node in another AST. %% The AST is searched in depth first order. -spec find_in_asts(OldNode, OldAst, NewAst) -> {found, NewNode} | not_found when OldNode :: erl_parse_tree(), OldAst :: erl_parse_tree() | [erl_parse_tree()], NewAst :: erl_parse_tree() | [erl_parse_tree()], NewNode :: erl_parse_tree(). find_in_asts(_OldNode, [], []) -> not_found; find_in_asts(OldNode, [X|Xs], [Y|Ys]) -> case find_in_asts(OldNode, X, Y) of {found, _} = Found -> Found; not_found -> find_in_asts(OldNode, Xs, Ys) end; find_in_asts(OldNode, OldNode, NewNode) -> {found, NewNode}; % Found the node! find_in_asts(OldNode, OldAst, NewAst) when is_tuple(OldAst), is_tuple(NewAst) -> find_in_asts(OldNode, gradualizer_lib:get_ast_children(OldAst), gradualizer_lib:get_ast_children(NewAst)); find_in_asts(_OldNode, _OldAst, _NewAst) -> not_found. % mismatch