Add format crate (#98)

Author: ferranbt

Cargo.lock

@@ -14,6 +14,7 @@ members = [
     "crates/pine-lexer",
     "crates/pine-ast",
     "crates/pine-diagnostics",
+    "crates/pine-format",
     "crates/pine-parser",
     "crates/pine-interpreter",
     "crates/pine-builtin-macro",
@@ -35,6 +36,7 @@ pine-lexer = { path = "crates/pine-lexer", version = "0.1.0" }
 pine-parser = { path = "crates/pine-parser", version = "0.1.0" }
 pine-ast = { path = "crates/pine-ast", version = "0.1.0" }
 pine-diagnostics = { path = "crates/pine-diagnostics", version = "0.1.0" }
+pine-format = { path = "crates/pine-format", version = "0.1.0" }
 pine-interpreter = { path = "crates/pine-interpreter", version = "0.1.0" }
 pine-builtin-macro = { path = "crates/pine-builtin-macro", version = "0.1.0" }
 pine-builtins = { path = "crates/pine-builtins", version = "0.1.0" }

Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "pine-format"
+version.workspace = true
+edition = "2021"
+description = "An opinionated formatter for Pine Script."
+license.workspace = true
+repository.workspace = true
+homepage.workspace = true
+
+[[bin]]
+name = "pine-format"
+path = "bin/pine-format.rs"
+
+[dependencies]
+pine-ast = { workspace = true }
+pine-core = { workspace = true }
+pine-lexer = { workspace = true }
+pine-parser = { workspace = true }
+clap = { workspace = true }

crates/pine-format/Cargo.toml

@@ -0,0 +1,56 @@
+use clap::Parser;
+use std::fs;
+use std::io::{self, Read, Write};
+
+#[derive(Parser)]
+#[command(name = "pine-format")]
+#[command(about = "Format PineScript source", long_about = None)]
+struct Cli {
+    /// Input file to format. Reads from stdin when omitted or with --stdin.
+    #[arg(value_name = "FILE")]
+    file: Option<String>,
+
+    /// Read input from stdin instead of a file.
+    #[arg(long)]
+    stdin: bool,
+
+    /// Rewrite the file in place instead of printing to stdout.
+    #[arg(short, long, requires = "file")]
+    write: bool,
+}
+
+fn main() {
+    let cli = Cli::parse();
+
+    let input = match &cli.file {
+        Some(filename) if !cli.stdin => fs::read_to_string(filename).unwrap_or_else(|e| {
+            eprintln!("Error reading file '{filename}': {e}");
+            std::process::exit(1);
+        }),
+        _ => {
+            let mut buffer = String::new();
+            io::stdin()
+                .read_to_string(&mut buffer)
+                .expect("Failed to read from stdin");
+            buffer
+        }
+    };
+
+    let formatted = match pine_format::format(&input) {
+        Ok(out) => out,
+        Err(e) => {
+            eprintln!("Format error: {e}");
+            std::process::exit(1);
+        }
+    };
+
+    if cli.write {
+        let filename = cli.file.expect("--write requires a file");
+        fs::write(&filename, formatted).unwrap_or_else(|e| {
+            eprintln!("Error writing file '{filename}': {e}");
+            std::process::exit(1);
+        });
+    } else {
+        io::stdout().write_all(formatted.as_bytes()).ok();
+    }
+}

crates/pine-format/bin/pine-format.rs

@@ -0,0 +1,77 @@
+use std::collections::{HashMap, VecDeque};
+
+use pine_lexer::{Token, TokenType};
+
+/// A piece of leading trivia above a statement, in source order.
+pub(crate) enum Lead {
+    Comment(String),
+    Blank,
+}
+
+/// The trivia recovered from the token stream: whole-line comments and blank
+/// lines (anchored before the statement that follows them) and trailing
+/// comments (kept on the code line they sit on).
+pub(crate) struct Comments {
+    leading: VecDeque<(u32, Lead)>,
+    trailing: HashMap<u32, String>,
+}
+
+impl Comments {
+    pub(crate) fn extract(tokens: &[Token]) -> Self {
+        let mut leading = VecDeque::new();
+        let mut trailing = HashMap::new();
+        let mut last_code_line: Option<usize> = None;
+
+        for token in tokens {
+            match &token.typ {
+                TokenType::Comment(text) => {
+                    let rendered = render(text);
+                    if last_code_line == Some(token.line) {
+                        trailing.insert(token.line as u32, rendered);
+                    } else {
+                        leading.push_back((token.line as u32, Lead::Comment(rendered)));
+                    }
+                }
+                TokenType::BlankLine => leading.push_back((token.line as u32, Lead::Blank)),
+                TokenType::Newline | TokenType::Indent | TokenType::Dedent | TokenType::Eof => {}
+                _ => last_code_line = Some(token.line),
+            }
+        }
+        Self { leading, trailing }
+    }
+
+    /// The comments and blank lines that sit above the statement starting on
+    /// `line`, in source order.
+    pub(crate) fn take_leading(&mut self, line: Option<u32>) -> Vec<Lead> {
+        let mut out = Vec::new();
+        let Some(line) = line else { return out };
+        while let Some((trivia_line, _)) = self.leading.front() {
+            if *trivia_line < line {
+                out.push(self.leading.pop_front().expect("front exists").1);
+            } else {
+                break;
+            }
+        }
+        out
+    }
+
+    /// The comment trailing the code on `line`, if any.
+    pub(crate) fn take_trailing(&mut self, line: Option<u32>) -> Option<String> {
+        line.and_then(|line| self.trailing.remove(&line))
+    }
+
+    /// Trailing comments never anchored to a statement's code line, in source
+    /// order. (Leftover leading trivia is drained via [`Comments::take_leading`]
+    /// with an unbounded line.)
+    pub(crate) fn drain_trailing(&mut self) -> Vec<String> {
+        let mut rest: Vec<(u32, String)> = self.trailing.drain().collect();
+        rest.sort_by_key(|(line, _)| *line);
+        rest.into_iter().map(|(_, text)| text).collect()
+    }
+}
+
+/// A comment's rendered form: `//` plus its verbatim text, trailing space
+/// trimmed. Verbatim keeps annotations (`//@version=5`) and dividers intact.
+fn render(text: &str) -> String {
+    format!("//{}", text.trim_end())
+}

crates/pine-format/src/comments.rs

@@ -0,0 +1,158 @@
+//! A Wadler/Prettier-style document algebra and its layout engine.
+//!
+//! Rules build a [`Doc`]; [`layout`] renders it, choosing for each [`Doc::Group`]
+//! whether it fits the target width flat or must break onto multiple lines.
+
+#[derive(Clone)]
+pub(crate) enum Doc {
+    Nil,
+    Text(String),
+    /// A space when flat, a newline (+ indent) when broken.
+    Line,
+    /// Nothing when flat, a newline (+ indent) when broken.
+    Softline,
+    /// Always a newline; forces every enclosing group to break.
+    Hardline,
+    Concat(Vec<Doc>),
+    Nest(usize, Box<Doc>),
+    Group(Box<Doc>),
+}
+
+pub(crate) fn text(s: impl Into<String>) -> Doc {
+    Doc::Text(s.into())
+}
+
+pub(crate) fn concat(docs: Vec<Doc>) -> Doc {
+    Doc::Concat(docs)
+}
+
+pub(crate) fn nest(indent: usize, doc: Doc) -> Doc {
+    Doc::Nest(indent, Box::new(doc))
+}
+
+pub(crate) fn group(doc: Doc) -> Doc {
+    Doc::Group(Box::new(doc))
+}
+
+pub(crate) fn line() -> Doc {
+    Doc::Line
+}
+
+pub(crate) fn softline() -> Doc {
+    Doc::Softline
+}
+
+pub(crate) fn hardline() -> Doc {
+    Doc::Hardline
+}
+
+/// `items` interleaved with `sep`.
+pub(crate) fn join(sep: Doc, items: Vec<Doc>) -> Doc {
+    let mut out = Vec::with_capacity(items.len() * 2);
+    for (i, doc) in items.into_iter().enumerate() {
+        if i > 0 {
+            out.push(sep.clone());
+        }
+        out.push(doc);
+    }
+    Doc::Concat(out)
+}
+
+#[derive(Clone, Copy, PartialEq)]
+enum Mode {
+    Flat,
+    Break,
+}
+
+pub(crate) fn layout(doc: &Doc, width: usize) -> String {
+    let mut out = String::new();
+    let mut col = 0usize;
+    let mut stack: Vec<(usize, Mode, &Doc)> = vec![(0, Mode::Break, doc)];
+
+    while let Some((indent, mode, doc)) = stack.pop() {
+        match doc {
+            Doc::Nil => {}
+            Doc::Text(s) => {
+                out.push_str(s);
+                col += s.chars().count();
+            }
+            Doc::Concat(docs) => {
+                for child in docs.iter().rev() {
+                    stack.push((indent, mode, child));
+                }
+            }
+            Doc::Nest(n, inner) => stack.push((indent + n, mode, inner)),
+            Doc::Line => match mode {
+                Mode::Flat => {
+                    out.push(' ');
+                    col += 1;
+                }
+                Mode::Break => col = newline(&mut out, indent),
+            },
+            Doc::Softline => match mode {
+                Mode::Flat => {}
+                Mode::Break => col = newline(&mut out, indent),
+            },
+            Doc::Hardline => col = newline(&mut out, indent),
+            Doc::Group(inner) => {
+                let mode = if fits(width.saturating_sub(col), indent, inner, &stack) {
+                    Mode::Flat
+                } else {
+                    Mode::Break
+                };
+                stack.push((indent, mode, inner));
+            }
+        }
+    }
+    out
+}
+
+fn newline(out: &mut String, indent: usize) -> usize {
+    out.push('\n');
+    for _ in 0..indent {
+        out.push(' ');
+    }
+    indent
+}
+
+/// Whether `group_inner` laid out flat, followed by whatever is already queued,
+/// stays within `remaining` columns before the next line break.
+fn fits<'a>(
+    remaining: usize,
+    indent: usize,
+    group_inner: &'a Doc,
+    rest: &[(usize, Mode, &'a Doc)],
+) -> bool {
+    let mut budget = remaining.min(isize::MAX as usize) as isize;
+    let mut stack: Vec<(usize, Mode, &Doc)> = rest.to_vec();
+    stack.push((indent, Mode::Flat, group_inner));
+
+    while budget >= 0 {
+        let (indent, mode, doc) = match stack.pop() {
+            Some(item) => item,
+            None => return true,
+        };
+        match doc {
+            Doc::Nil => {}
+            Doc::Text(s) => budget -= s.chars().count() as isize,
+            Doc::Concat(docs) => {
+                for child in docs.iter().rev() {
+                    stack.push((indent, mode, child));
+                }
+            }
+            Doc::Nest(n, inner) => stack.push((indent + n, mode, inner)),
+            Doc::Line => match mode {
+                Mode::Flat => budget -= 1,
+                Mode::Break => return true,
+            },
+            Doc::Softline => match mode {
+                Mode::Flat => {}
+                Mode::Break => return true,
+            },
+            // A forced break ends this line, so everything up to here fit.
+            Doc::Hardline => return true,
+            Doc::Group(inner) => stack.push((indent, Mode::Flat, inner)),
+        }
+    }
+    false
+}