Added layout.highlighters. A new, much faster way to do selection and search highlighting.

6 files changed, 391 insertions, 46 deletions

diff --git a/prompt_toolkit/layout/containers.py b/prompt_toolkit/layout/containers.py
index d536b47c..1a664383 100644
--- a/prompt_toolkit/layout/containers.py
+++ b/prompt_toolkit/layout/containers.py
@@ -4,11 +4,12 @@ Container for the layout.
 """
 from __future__ import unicode_literals
 
-from six import with_metaclass
 from abc import ABCMeta, abstractmethod
+from collections import defaultdict
 from pygments.token import Token
+from six import with_metaclass
 
-from .screen import Point, WritePosition
+from .screen import Point, WritePosition, Char
 from .dimension import LayoutDimension, sum_layout_dimensions, max_layout_dimensions
 from .controls import UIControl, TokenListControl
 from .margins import Margin
@@ -868,21 +869,26 @@ class Window(Container):
         Write window to screen. This renders the user control, the margins and
         copies everything over to the absolute position at the given screen.
         """
-        # Render margins.
+        # Calculate margin sizes.
         left_margin_widths = [m.get_width(cli) for m in self.left_margins]
         right_margin_widths = [m.get_width(cli) for m in self.right_margins]
         total_margin_width = sum(left_margin_widths + right_margin_widths)
 
         # Render UserControl.
-        temp_screen = self.content.create_screen(
+        tpl = self.content.create_screen(
             cli, write_position.width - total_margin_width, write_position.height)
+        if isinstance(tpl, tuple):
+            temp_screen, highlighting = tpl
+        else:
+            # For backwards, compatibility.
+            temp_screen, highlighting = tpl, defaultdict(lambda: defaultdict(lambda: None))
 
         # Scroll content.
         applied_scroll_offsets = self._scroll(
             temp_screen, write_position.width - total_margin_width, write_position.height, cli)
 
         # Write body to screen.
-        self._copy_body(cli, temp_screen, screen, write_position,
+        self._copy_body(cli, temp_screen, highlighting, screen, write_position,
                         sum(left_margin_widths), write_position.width - total_margin_width,
                         applied_scroll_offsets)
 
@@ -959,7 +965,8 @@ class Window(Container):
             self._copy_margin(margin_screen, screen, write_position, move_x, width)
             move_x += width
 
-    def _copy_body(self, cli, temp_screen, new_screen, write_position, move_x, width, applied_scroll_offsets):
+    def _copy_body(self, cli, temp_screen, highlighting, new_screen,
+                   write_position, move_x, width, applied_scroll_offsets):
         """
         Copy characters from the temp screen that we got from the `UIControl`
         to the real screen.
@@ -971,6 +978,9 @@ class Window(Container):
         temp_buffer = temp_screen.data_buffer
         new_buffer = new_screen.data_buffer
         temp_screen_height = temp_screen.height
+
+        vertical_scroll = self.vertical_scroll
+        horizontal_scroll = self.horizontal_scroll
         y = 0
 
         # Now copy the region we need to the real screen.
@@ -985,25 +995,31 @@ class Window(Container):
                 # screen's height.)
                 break
             else:
-                temp_row = temp_buffer[y + self.vertical_scroll]
+                temp_row = temp_buffer[y + vertical_scroll]
+                highlighting_row = highlighting[y + vertical_scroll]
 
                 # Copy row content, except for transparent tokens.
                 # (This is useful in case of floats.)
+                # Also apply highlighting.
                 for x in range(0, width):
-                    cell = temp_row[x + self.horizontal_scroll]
-                    if cell.token != Transparent:
+                    cell = temp_row[x + horizontal_scroll]
+                    highlighting_token = highlighting_row[x]
+
+                    if highlighting_token:
+                        new_row[x + xpos] = Char(cell.char, highlighting_token)
+                    elif cell.token != Transparent:
                         new_row[x + xpos] = cell
 
         if self.content.has_focus(cli):
-            new_screen.cursor_position = Point(y=temp_screen.cursor_position.y + ypos - self.vertical_scroll,
-                                               x=temp_screen.cursor_position.x + xpos - self.horizontal_scroll)
+            new_screen.cursor_position = Point(y=temp_screen.cursor_position.y + ypos - vertical_scroll,
+                                               x=temp_screen.cursor_position.x + xpos - horizontal_scroll)
 
             if not self.always_hide_cursor(cli):
                 new_screen.show_cursor = temp_screen.show_cursor
 
         if not new_screen.menu_position and temp_screen.menu_position:
-            new_screen.menu_position = Point(y=temp_screen.menu_position.y + ypos - self.vertical_scroll,
-                                             x=temp_screen.menu_position.x + xpos - self.horizontal_scroll)
+            new_screen.menu_position = Point(y=temp_screen.menu_position.y + ypos - vertical_scroll,
+                                             x=temp_screen.menu_position.x + xpos - horizontal_scroll)
 
         # Update height of the output screen. (new_screen.write_data is not
         # called, so the screen is not aware of its height.)
diff --git a/prompt_toolkit/layout/controls.py b/prompt_toolkit/layout/controls.py
index 52e2d881..f131142a 100644
--- a/prompt_toolkit/layout/controls.py
+++ b/prompt_toolkit/layout/controls.py
@@ -4,8 +4,9 @@ User interface Controls for the layout.
 from __future__ import unicode_literals
 from pygments.token import Token
 
-from six import with_metaclass
 from abc import ABCMeta, abstractmethod
+from collections import defaultdict, namedtuple
+from six import with_metaclass
 
 from prompt_toolkit.enums import DEFAULT_BUFFER
 from prompt_toolkit.filters import to_cli_filter
@@ -14,6 +15,7 @@ from prompt_toolkit.search_state import SearchState
 from prompt_toolkit.selection import SelectionType
 from prompt_toolkit.utils import get_cwidth, SimpleLRUCache
 
+from .highlighters import Highlighter
 from .lexers import Lexer, SimpleLexer
 from .processors import Processor, Transformation
 from .screen import Screen, Char, Point
@@ -58,8 +60,13 @@ class UIControl(with_metaclass(ABCMeta, object)):
     def create_screen(self, cli, width, height):
         """
         Write the content at this position to the screen.
+
+        Returns a :class:`.Screen` instance.
+
+        Optionally, this can also return a (screen, highlighting) tuple, where
+        the `highlighting` is a dictionary of dictionaries. Mapping
+        y->x->Token if this position needs to be highlighted with that Token.
         """
-        pass
 
     def mouse_handler(self, cli, mouse_event):
         """
@@ -220,7 +227,9 @@ class TokenListControl(UIControl):
                     tokens2.extend(process_line(line))
                 tokens = tokens2
 
-            indexes_to_pos = screen.write_data(tokens, width=(width if wrap_lines else None))
+            write_data_result = screen.write_data(tokens, width=(width if wrap_lines else None))
+
+            indexes_to_pos = write_data_result.indexes_to_pos
             pos_to_indexes = dict((v, k) for k, v in indexes_to_pos.items())
         else:
             pos_to_indexes = {}
@@ -306,6 +315,7 @@ class BufferControl(UIControl):
     def __init__(self,
                  buffer_name=DEFAULT_BUFFER,
                  input_processors=None,
+                 highlighters=None,
                  lexer=None,
                  preview_search=False,
                  wrap_lines=True,
@@ -313,6 +323,7 @@ class BufferControl(UIControl):
                  default_char=None,
                  focus_on_click=False):
         assert input_processors is None or all(isinstance(i, Processor) for i in input_processors)
+        assert highlighters is None or all(isinstance(i, Highlighter) for i in highlighters)
         assert menu_position is None or callable(menu_position)
         assert lexer is None or isinstance(lexer, Lexer)
 
@@ -321,6 +332,7 @@ class BufferControl(UIControl):
         self.focus_on_click = to_cli_filter(focus_on_click)
 
         self.input_processors = input_processors or []
+        self.highlighters = highlighters or []
         self.buffer_name = buffer_name
         self.menu_position = menu_position
         self.lexer = lexer or SimpleLexer()
@@ -337,6 +349,11 @@ class BufferControl(UIControl):
         #: to recreate the same screen again.)
         self._screen_lru_cache = SimpleLRUCache(maxsize=8)
 
+        #: Highlight Cache.
+        #: When nothing of the buffer content or processors has changed, but
+        #: the highlighting of the selection/search changes,
+        self._highlight_lru_cache = SimpleLRUCache(maxsize=8)
+
         self._xy_to_cursor_position = None
         self._last_click_timestamp = None
 
@@ -362,7 +379,7 @@ class BufferControl(UIControl):
     def preferred_height(self, cli, width):
         # Draw content on a screen using this width. Measure the height of the
         # result.
-        screen = self.create_screen(cli, width, None)
+        screen, highlighters = self.create_screen(cli, width, None)
         return screen.height
 
     def _get_input_tokens(self, cli, document):
@@ -457,7 +474,10 @@ class BufferControl(UIControl):
             input_tokens, source_to_display, display_to_source = self._get_input_tokens(cli, document)
             input_tokens += [(self.default_char.token, ' ')]
 
-            indexes_to_pos = screen.write_data(input_tokens, width=wrap_width)
+            write_data_result = screen.write_data(input_tokens, width=wrap_width)
+            indexes_to_pos = write_data_result.indexes_to_pos
+            line_lengths = write_data_result.line_lengths
+
             pos_to_indexes = dict((v, k) for k, v in indexes_to_pos.items())
 
             def cursor_position_to_xy(cursor_position):
@@ -495,7 +515,7 @@ class BufferControl(UIControl):
                 # Transform.
                 return display_to_source(index)
 
-            return screen, cursor_position_to_xy, xy_to_cursor_position
+            return screen, cursor_position_to_xy, xy_to_cursor_position, line_lengths
 
         # Build a key for the caching. If any of these parameters changes, we
         # have to recreate a new screen.
@@ -512,7 +532,8 @@ class BufferControl(UIControl):
         )
 
         # Get from cache, or create if this doesn't exist yet.
-        screen, cursor_position_to_xy, self._xy_to_cursor_position = self._screen_lru_cache.get(key, _create_screen)
+        screen, cursor_position_to_xy, self._xy_to_cursor_position, line_lengths = \
+            self._screen_lru_cache.get(key, _create_screen)
 
         x, y = cursor_position_to_xy(document.cursor_position)
         screen.cursor_position = Point(y=y, x=x)
@@ -539,7 +560,63 @@ class BufferControl(UIControl):
             else:
                 screen.menu_position = None
 
-        return screen
+        # Add highlighting.
+        highlight_key = (
+            key,  # Includes everything from the 'key' above. (E.g. when the
+                     # document changes, we have to recalculate highlighting.)
+
+            # Include invalidation_hashes from all highlighters.
+            tuple(h.invalidation_hash(cli, document) for h in self.highlighters)
+        )
+
+        highlighting = self._highlight_lru_cache.get(highlight_key, lambda:
+            self._get_highlighting(cli, document, cursor_position_to_xy, line_lengths))
+
+        return screen, highlighting
+
+    def _get_highlighting(self, cli, document, cursor_position_to_xy, line_lengths):
+        """
+        Return a _HighlightDict for the highlighting. (This is a lazy dict of dicts.)
+
+        The Window class will apply this for the visible regions. - That way,
+        we don't have to recalculate the screen again for each selection/search
+        change.
+
+        :param line_lengths: Maps line numbers to the length of these lines.
+        """
+        def get_row_size(y):
+            " Return the max 'x' value for a given row in the screen. "
+            return max(1, line_lengths.get(y, 0))
+
+        # Get list of fragments.
+        row_to_fragments = defaultdict(list)
+
+        for h in self.highlighters:
+            for fragment in h.get_fragments(cli, document):
+                # Expand fragments.
+                start_column, start_row = cursor_position_to_xy(fragment.start)
+                end_column, end_row = cursor_position_to_xy(fragment.end)
+                token = fragment.token
+
+                if start_row == end_row:
+                    # Single line highlighting.
+                    row_to_fragments[start_row].append(
+                        _HighlightFragment(start_column, end_column, token))
+                else:
+                    # Multi line highlighting.
+                    # (First line.)
+                    row_to_fragments[start_row].append(
+                        _HighlightFragment(start_column, get_row_size(start_row), token))
+
+                    # (Middle lines.)
+                    for y in range(start_row + 1, end_row):
+                        row_to_fragments[y].append(_HighlightFragment(0, get_row_size(y), token))
+
+                    # (Last line.)
+                    row_to_fragments[end_row].append(_HighlightFragment(0, end_column, token))
+
+        # Create dict to return.
+        return _HighlightDict(row_to_fragments)
 
     def mouse_handler(self, cli, mouse_event):
         """
@@ -602,3 +679,41 @@ class BufferControl(UIControl):
     def move_cursor_up(self, cli):
         b = self._buffer(cli)
         b.cursor_position += b.document.get_cursor_up_position()
+
+
+_HighlightFragment = namedtuple('_HighlightFragment', 'start_column end_column token')  # End is excluded.
+
+
+class _HighlightDict(dict):
+    """
+    Helper class to contain the highlighting.
+    Maps 'y' coordinate to 'x' coordinate to Token.
+
+    :param row_to_fragments: Dictionary that maps row numbers to list of `_HighlightFragment`.
+    """
+    def __init__(self, row_to_fragments):
+        self.row_to_fragments = row_to_fragments
+
+    def __missing__(self, key):
+        result = _HighlightDictRow(self.row_to_fragments[key])
+        self[key] = result
+        return result
+
+    def __repr__(self):
+        return '_HighlightDict(%r)' % (dict.__repr__(self), )
+
+
+class _HighlightDictRow(dict):
+    def __init__(self, list_of_fragments):
+        self.list_of_fragments = list_of_fragments
+
+    def __missing__(self, key):
+        result = None
+
+        for f in self.list_of_fragments:
+            if f.start_column <= key < f.end_column:  # End is excluded.
+                result = f.token
+                break
+
+        self[key] = result
+        return result
diff --git a/prompt_toolkit/layout/highlighters.py b/prompt_toolkit/layout/highlighters.py
new file mode 100644
index 00000000..3934baa5
--- /dev/null
+++ b/prompt_toolkit/layout/highlighters.py
@@ -0,0 +1,204 @@
+"""
+Highlighters for usage in a BufferControl.
+
+Highlighters are very similar to processors, but they are applied after the
+BufferControl created a screen instance. (Instead of right before creating the screen.)
+Highlighters can't change the content of the screen, but they can mark regions
+(start_pos, end_pos) as highlighted, using a certain Token.
+
+When possible, it's adviced to use a Highlighter instead of a Processor,
+because most of the highlighting code is applied only to the visible region of
+the screen. (The Window class will apply the highlighting to the visible region.)
+"""
+from __future__ import unicode_literals
+from pygments.token import Token
+from abc import ABCMeta, abstractmethod
+from six import with_metaclass
+
+from prompt_toolkit.document import Document
+from prompt_toolkit.enums import SEARCH_BUFFER
+from prompt_toolkit.filters import to_cli_filter
+
+__all__ = (
+    'Fragment',
+    'SelectionHighlighter',
+    'SearchHighlighter',
+    'MatchingBracketHighlighter',
+    'ConditionalHighlighter',
+)
+
+
+class Fragment(object):
+    """
+    Highlight fragment.
+
+    :param start: (int) Cursor start position.
+    :param end: (int) Cursor end position.
+    :param token: Pygments Token.
+    """
+    def __init__(self, start, end, token):
+        self.start = start
+        self.end = end
+        self.token = token
+
+    def __repr__(self):
+        return 'Fragment(%r, %r, %r)' % (self.start, self.end, self.token)
+
+
+class Highlighter(with_metaclass(ABCMeta, object)):
+    @abstractmethod
+    def get_fragments(self, cli, document):
+        """
+        Return a list of :class:`.Fragment` instances.
+        (This can be a generator as well.)
+        """
+        return []
+
+
+class SelectionHighlighter(Highlighter):
+    """
+    Highlight the selection.
+    """
+    def get_fragments(self, cli, document):
+        selection_range = document.selection_range()
+
+        if selection_range:
+            from_, to = selection_range
+            yield Fragment(from_, to, Token.SelectedText)
+
+    def invalidation_hash(self, cli, document):
+        # When the selection range changes, highlighting will be different.
+        return (
+            document.selection_range(),
+        )
+
+
+class SearchHighlighter(Highlighter):
+    """
+    Highlight search matches in the document.
+
+    :param preview_search: A Filter; when active it indicates that we take
+        the search text in real time while the user is typing, instead of the
+        last active search state.
+    """
+    def __init__(self, preview_search=False, search_buffer_name=SEARCH_BUFFER):
+        self.preview_search = to_cli_filter(preview_search)
+        self.search_buffer_name = search_buffer_name
+
+    def _get_search_text(self, cli):
+        """
+        The text we are searching for.
+        """
+        # When the search buffer has focus, take that text.
+        if self.preview_search(cli) and cli.buffers[self.search_buffer_name].text:
+            return cli.buffers[self.search_buffer_name].text
+        # Otherwise, take the text of the last active search.
+        else:
+            return cli.search_state.text
+
+    def get_fragments(self, cli, document):
+        search_text = self._get_search_text(cli)
+        ignore_case = cli.is_ignoring_case
+
+        if search_text and not cli.is_returning:
+            for index in document.find_all(search_text, ignore_case=ignore_case):
+                if index == document.cursor_position:
+                    token = Token.SearchMatch.Current
+                else:
+                    token = Token.SearchMatch
+
+                yield Fragment(index, index + len(search_text), token)
+
+    def invalidation_hash(self, cli, document):
+        search_text = self._get_search_text(cli)
+
+        # When the search state changes, highlighting will be different.
+        return (
+            search_text,
+            cli.is_returning,
+
+            # When we search for text, and the cursor position changes. The
+            # processor has to be applied every time again, because the current
+            # match is highlighted in another color.
+            (search_text and document.cursor_position),
+        )
+
+
+class ConditionalHighlighter(Highlighter):
+    """
+    Highlighter that applies another highlighter, according to a certain condition.
+
+    :param highlighter: :class:`.Highlighter` instance.
+    :param filter: :class:`~prompt_toolkit.filters.CLIFilter` instance.
+    """
+    def __init__(self, highlighter, filter):
+        assert isinstance(highlighter, Highlighter)
+
+        self.highlighter = highlighter
+        self.filter = to_cli_filter(filter)
+
+    def get_fragments(self, cli, document):
+        if self.filter(cli):
+            return self.highlighter.get_fragments(cli, document)
+        else:
+            return []
+
+    def invalidation_hash(self, cli, document):
+        # When enabled, use the hash of the highlighter. Otherwise, just use
+        # False.
+        if self.filter(cli):
+            return (True, self.highlighter.invalidation_hash(cli, document))
+        else:
+            return False
+
+
+class MatchingBracketHighlighter(Highlighter):
+    """
+    When the cursor is on or right after a bracket, it highlights the matching
+    bracket.
+    """
+    _closing_braces = '])}>'
+
+    def __init__(self, chars='[](){}<>'):
+        self.chars = chars
+
+    def get_fragments(self, cli, document):
+        result = []
+
+        def replace_token(pos):
+            """ Replace token in list of tokens. """
+            result.append(Fragment(pos, pos + 1, Token.MatchingBracket))
+
+        def apply_for_document(document):
+            """ Find and replace matching tokens. """
+            if document.current_char in self.chars:
+                pos = document.matching_bracket_position
+
+                if pos:
+                    replace_token(document.cursor_position)
+                    replace_token(document.cursor_position + pos)
+                    return True
+
+        # Apply for character below cursor.
+        applied = apply_for_document(document)
+
+        # Otherwise, apply for character before cursor.
+        if (not applied and document.cursor_position > 0 and
+                document.char_before_cursor in self._closing_braces):
+            apply_for_document(Document(document.text, document.cursor_position - 1))
+
+        return result
+
+    def invalidation_hash(self, cli, document):
+        on_brace = document.current_char in self.chars
+        after_brace = document.char_before_cursor in self.chars
+
+        if on_brace:
+            return (True, document.cursor_position)
+        elif after_brace and document.char_before_cursor in self._closing_braces:
+            return (True, document.cursor_position - 1)
+        else:
+            # Don't include the cursor position in the hash if we are not *on*
+            # a brace. We don't have to rerender the output, because it will be
+            # the same anyway.
+            return False
diff --git a/prompt_toolkit/layout/processors.py b/prompt_toolkit/layout/processors.py
index f15935fd..0a0b3a03 100644
--- a/prompt_toolkit/layout/processors.py
+++ b/prompt_toolkit/layout/processors.py
@@ -87,7 +87,7 @@ class Transformation(object):
         self.display_to_source = display_to_source or (lambda i: i)
 
 
-class HighlightSearchProcessor(Processor):
+class HighlightSearchProcessor(Processor):  # XXX: Deprecated!
     """
     Processor that highlights search matches in the document.
 
@@ -95,16 +95,17 @@ class HighlightSearchProcessor(Processor):
         the search text in real time while the user is typing, instead of the
         last active search state.
     """
-    def __init__(self, preview_search=False):
+    def __init__(self, preview_search=False, search_buffer_name=SEARCH_BUFFER):
         self.preview_search = to_cli_filter(preview_search)
+        self.search_buffer_name = search_buffer_name
 
     def _get_search_text(self, cli):
         """
         The text we are searching for.
         """
         # When the search buffer has focus, take that text.
-        if self.preview_search(cli) and cli.is_searching and cli.buffers[SEARCH_BUFFER].text:
-            return cli.buffers[SEARCH_BUFFER].text
+        if self.preview_search(cli) and cli.buffers[self.search_buffer_name].text:
+            return cli.buffers[self.search_buffer_name].text
         # Otherwise, take the text of the last active search.
         else:
             return cli.search_state.text
@@ -143,7 +144,7 @@ class HighlightSearchProcessor(Processor):
         )
 
 
-class HighlightSelectionProcessor(Processor):
+class HighlightSelectionProcessor(Processor):  # XXX: Deprecated!
     """
     Processor that highlights the selection in the document.
     """
@@ -182,7 +183,7 @@ class PasswordProcessor(Processor):
             document, [(token, self.char * len(text)) for token, text in tokens])
 
 
-class HighlightMatchingBracketProcessor(Processor):
+class HighlightMatchingBracketProcessor(Processor):  # XXX: Deprecated!
     """
     When the cursor is on or right after a bracket, it highlights the matching
     bracket.
diff --git a/prompt_toolkit/layout/screen.py b/prompt_toolkit/layout/screen.py
index 7edf0f94..57dc2050 100644
--- a/prompt_toolkit/layout/screen.py
+++ b/prompt_toolkit/layout/screen.py
@@ -102,6 +102,8 @@ class CharCache(dict):
 _CHAR_CACHE = CharCache()
 Transparent = Token.Transparent
 
+WriteDataResult = namedtuple('WriteDataResult', 'indexes_to_pos line_lengths')
+
 
 class Screen(object):
     """
@@ -163,8 +165,8 @@ class Screen(object):
         line_number = 0
         requires_line_feed = True
         indexes_to_pos = {}  # Map input positions to (x, y) coordinates.
+        line_lengths = {}  # Map line numbers (y) to max_x for this line.
         set_cursor_position = Token.SetCursorPosition
-        max_x = 0
 
         for token, text in data:
             if token == set_cursor_position:
@@ -182,7 +184,7 @@ class Screen(object):
                 # In case there is no more place left at this line, go first to the
                 # following line. (Also in case of double-width characters.)
                 if x + char_width > max_allowed_x and char != '\n':
-                    max_x = max(max_x, x)
+                    line_lengths[y] = x
                     y += 1
                     x = 0
 
@@ -191,7 +193,7 @@ class Screen(object):
 
                 # Insertion of newline
                 if char == '\n':
-                    max_x = max(max_x, x)
+                    line_lengths[y] = x
                     y += 1
                     x = 0
                     requires_line_feed = True
@@ -214,10 +216,12 @@ class Screen(object):
 
                 index += 1
 
+        line_lengths[y] = x
+
         self.height = max(self.height, y + 1)
-        self.width = max(self.width, max_x, x)
+        self.width = max(self.width, max(line_lengths.values()))
 
-        return indexes_to_pos
+        return WriteDataResult(indexes_to_pos, line_lengths)
 
     def replace_all_tokens(self, token):
         """
diff --git a/prompt_toolkit/shortcuts.py b/prompt_toolkit/shortcuts.py
index b3f40e2b..40ee2aa7 100644
--- a/prompt_toolkit/shortcuts.py
+++ b/prompt_toolkit/shortcuts.py
@@ -34,7 +34,8 @@ from .layout.controls import BufferControl, TokenListControl
 from .layout.dimension import LayoutDimension
 from .layout.lexers import PygmentsLexer
 from .layout.menus import CompletionsMenu, MultiColumnCompletionsMenu
-from .layout.processors import PasswordProcessor, HighlightSearchProcessor, HighlightSelectionProcessor, ConditionalProcessor, AppendAutoSuggestion
+from .layout.processors import PasswordProcessor, ConditionalProcessor, AppendAutoSuggestion
+from .layout.highlighters import SearchHighlighter, SelectionHighlighter, ConditionalHighlighter
 from .layout.prompt import DefaultPrompt
 from .layout.screen import Char
 from .layout.toolbars import ValidationToolbar, SystemToolbar, ArgToolbar, SearchToolbar
@@ -196,19 +197,21 @@ def create_prompt_layout(message='', lexer=None, is_password=False,
     except TypeError: # Happens when lexer is `None` or an instance of something else.
         pass
 
-    # Create processors list.
-    # (DefaultPrompt should always be at the end.)
-    input_processors = [ConditionalProcessor(
-                            # By default, only highlight search when the search
-                            # input has the focus. (Note that this doesn't mean
-                            # there is no search: the Vi 'n' binding for instance
-                            # still allows to jump to the next match in
-                            # navigation mode.)
-                            HighlightSearchProcessor(preview_search=True),
-                            HasFocus(SEARCH_BUFFER)),
-                        HighlightSelectionProcessor(),
-                        ConditionalProcessor(AppendAutoSuggestion(), HasFocus(DEFAULT_BUFFER) & ~IsDone()),
-                        ConditionalProcessor(PasswordProcessor(), is_password)]
+    # Create highlighters and processors list.
+    highlighters = [
+        ConditionalHighlighter(
+            # By default, only highlight search when the search
+            # input has the focus. (Note that this doesn't mean
+            # there is no search: the Vi 'n' binding for instance
+            # still allows to jump to the next match in
+            # navigation mode.)
+            SearchHighlighter(preview_search=True),
+        HasFocus(SEARCH_BUFFER)),
+        SelectionHighlighter()]
+
+    input_processors = [
+        ConditionalProcessor(AppendAutoSuggestion(), HasFocus(DEFAULT_BUFFER) & ~IsDone()),
+        ConditionalProcessor(PasswordProcessor(), is_password)]
 
     if extra_input_processors:
         input_processors.extend(extra_input_processors)
@@ -216,6 +219,7 @@ def create_prompt_layout(message='', lexer=None, is_password=False,
     # Show the prompt before the input (using the DefaultPrompt processor.
     # This also replaces it with reverse-i-search and 'arg' when required.
     # (Only for single line mode.)
+    # (DefaultPrompt should always be at the end of the processors.)
     input_processors.append(ConditionalProcessor(
         DefaultPrompt(get_prompt_tokens), ~multiline))
 
@@ -258,6 +262,7 @@ def create_prompt_layout(message='', lexer=None, is_password=False,
             FloatContainer(
                 Window(
                     BufferControl(
+                        highlighters=highlighters,
                         input_processors=input_processors,
                         lexer=lexer,
                         wrap_lines=wrap_lines,