From dc081151bcd029d49aaf00a629653627fdc99aca Mon Sep 17 00:00:00 2001
From: Jonathan Turner <jonathan.d.turner@gmail.com>
Date: Sun, 23 Jun 2019 07:47:29 +1200
Subject: [PATCH] add back some shell coloring

---
 src/parser/parse2/token_tree.rs |  4 +-
 src/shell/helper.rs             | 70 ++++++++++++++++++++++++++-------
 2 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/src/parser/parse2/token_tree.rs b/src/parser/parse2/token_tree.rs
index e618feaa93..d2c6e7ffb8 100644
--- a/src/parser/parse2/token_tree.rs
+++ b/src/parser/parse2/token_tree.rs
@@ -94,10 +94,10 @@ pub struct PathNode {
 
 #[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Getters, new)]
 pub struct PipelineElement {
-    pre_ws: Option<Span>,
+    pub pre_ws: Option<Span>,
     #[get = "crate"]
     call: Spanned<CallNode>,
-    post_ws: Option<Span>,
+    pub post_ws: Option<Span>,
 }
 
 impl PipelineElement {
diff --git a/src/shell/helper.rs b/src/shell/helper.rs
index 85ad50de6a..d91f55f0c0 100644
--- a/src/shell/helper.rs
+++ b/src/shell/helper.rs
@@ -1,5 +1,8 @@
 use crate::shell::completer::NuCompleter;
-
+use crate::parser::parse2::PipelineElement;
+use crate::parser::parse2::token_tree::TokenNode;
+use crate::parser::parse2::tokens::RawToken;
+use crate::parser::parse2::span::Spanned;
 use crate::parser::nom_input;
 use crate::prelude::*;
 use ansi_term::Color;
@@ -68,26 +71,24 @@ impl Highlighter for Helper {
 
                 let mut iter = tokens.into_iter();
 
-                match iter.next() {
-                    None => return Cow::Owned(out),
-                    Some(v) => out.push_str(v.span().slice(line)),
-                };
+                // match iter.next() {
+                //     None => return Cow::Owned(out),
+                //     Some(v) => out.push_str(v.span().slice(line)),
+                // };
+                let mut first = true;
 
                 loop {
                     match iter.next() {
                         None => return Cow::Owned(out),
                         Some(token) => {
-                            // let styled = token_style(&token, state);
+                            let styled = paint_pipeline_element(&token, line);
 
-                            // trace!("token={:?}", token);
-                            // trace!("style={:?}", style);
-                            // trace!("new_state={:?}", new_state);
+                            if !first {
+                                out.push_str("|");
+                            } else {
+                                first = false;
+                            }
 
-                            // state = new_state;
-                            // let slice = &line[*start..*end];
-                            // let styled = style.paint(slice);
-                            out.push_str("|");
-                            let styled = Color::Black.bold().paint(token.span().slice(line));
                             out.push_str(&styled.to_string());
                         }
                     }
@@ -101,4 +102,45 @@ impl Highlighter for Helper {
     }
 }
 
+fn paint_token_node(token_node: &TokenNode, line: &str) -> String {
+    let styled = match token_node {
+        TokenNode::Call(..) => Color::Cyan.bold().paint(token_node.span().slice(line)),
+        TokenNode::Whitespace(..) => Color::White.normal().paint(token_node.span().slice(line)),
+        TokenNode::Flag(..) => Color::Black.bold().paint(token_node.span().slice(line)),
+        TokenNode::Identifier(..) => Color::Yellow.bold().paint(token_node.span().slice(line)),
+        TokenNode::Path(..) => Color::Green.bold().paint(token_node.span().slice(line)),
+        TokenNode::Error(..) => Color::Red.bold().paint(token_node.span().slice(line)),
+        TokenNode::Delimited(..) => Color::White.paint(token_node.span().slice(line)),
+        TokenNode::Operator(..) => Color::Purple.bold().paint(token_node.span().slice(line)),
+        TokenNode::Pipeline(..) => Color::Blue.normal().paint(token_node.span().slice(line)),
+        TokenNode::Token(Spanned { item: RawToken::Integer(..), ..}) => Color::Purple.bold().paint(token_node.span().slice(line)),
+        TokenNode::Token(Spanned { item: RawToken::Size(..), ..}) => Color::Purple.bold().paint(token_node.span().slice(line)),
+        TokenNode::Token(Spanned { item: RawToken::String(..), ..}) => Color::Green.normal().paint(token_node.span().slice(line)),
+        TokenNode::Token(Spanned { item: RawToken::Variable(..), ..}) => Color::Yellow.bold().paint(token_node.span().slice(line)),
+        TokenNode::Token(Spanned { item: RawToken::Bare, ..}) => Color::Green.normal().paint(token_node.span().slice(line)),
+    };
+
+    styled.to_string()
+}
+
+fn paint_pipeline_element(pipeline_element: &PipelineElement, line: &str) -> String {
+    let mut styled = String::new();
+
+    if let Some(ws) = pipeline_element.pre_ws {
+        styled.push_str(&Color::White.normal().paint(ws.slice(line)));
+    }
+    styled.push_str(&paint_token_node(pipeline_element.call().head(), line));
+
+    if let Some(children) = pipeline_element.call().children() {
+        for child in children {
+            styled.push_str(&paint_token_node(child, line));
+        }
+    }
+    if let Some(ws) = pipeline_element.post_ws {
+        styled.push_str(&Color::White.normal().paint(ws.slice(line)));
+    }
+
+    styled.to_string()
+}
+
 impl rustyline::Helper for Helper {}