blues-lsp 0.1.3

use std::{marker::PhantomData, str::CharIndices};

use arcstr::Substr;

use crate::util::ext::{
    arcstr::{EMPTY_SUBSTR, SubStrExt},
    option::OptionThenExt,
};

use super::{
    TextChunk,
    location::{LineColCounter, Pos},
};

pub trait CharStream<'src>: Clone {
    fn peek(&self) -> Option<char>;
    // Returns `true` if chunk boundary was crossed
    fn step(&mut self) -> bool;

    fn next(&mut self) -> Option<char> {
        let out = self.peek()?;
        self.step();
        Some(out)
    }

    fn next_peek(&mut self) -> Option<char> {
        self.next();
        self.peek()
    }

    fn char_idx(&self) -> usize;
    fn pos(&self) -> Pos;

    fn current_chunk(&self) -> &'src TextChunk;

    fn chunk_state<'slf>(&'slf mut self) -> ChunkState<'slf, 'src, Self> {
        ChunkState::new(self)
    }

    fn expect_any(&mut self) {
        self.next().expect("consume_any called on empty chunk");
    }

    fn expect(&mut self, expected: char) {
        assert_eq!(self.next(), Some(expected))
    }

    fn read_while(&mut self, mut pred: impl FnMut(char) -> bool) -> Substr {
        let mut s = EMPTY_SUBSTR;
        let mut sub = self.chunk_state().writer(|c| s.append(c.text));

        while let Some(c) = sub.peek() {
            if !pred(c) {
                break;
            }
            sub.advance();
        }
        sub.flush();
        s
    }

    fn skip_while(&mut self, mut pred: impl FnMut(char) -> bool) -> usize {
        let mut cnt = 0;

        while let Some(c) = self.peek() {
            if !pred(c) {
                break;
            }
            cnt += 1;
            self.expect_any();
        }

        cnt
    }
}

#[derive(Clone)]
pub struct TextCharStream<'src, I: Iterator<Item = &'src TextChunk>> {
    chunks: I,

    current_chunk: &'src TextChunk,
    current_chars: CharIndices<'src>,
    current_pos: LineColCounter,

    peeked: Option<(usize, char)>,
}

static EMPTY_CHUNK: TextChunk = TextChunk {
    pos: Pos::zero(),
    text: EMPTY_SUBSTR,
};

impl<'src, I: Iterator<Item = &'src TextChunk>> TextCharStream<'src, I> {
    pub fn new(input: I) -> Self {
        let mut stream = Self {
            chunks: input,

            // Start with empty iterator
            current_chunk: &EMPTY_CHUNK,
            current_chars: "".char_indices(),
            current_pos: LineColCounter::default(),

            peeked: None,
        };

        stream.peeked = stream.next_new();

        stream
    }

    /// Returns first char of new chunk, or None if no chunks are left
    fn next_new(&mut self) -> Option<(usize, char)> {
        let chunk = self.chunks.next()?;

        self.current_chunk = chunk;
        self.current_chars = chunk.text.char_indices();
        self.current_pos = LineColCounter::start_at(chunk.pos.text_pos);

        Some(
            self.current_chars
                .next()
                .expect("New chunks cannot be empty"),
        )
    }
}

impl<'src, I: Iterator<Item = &'src TextChunk> + Clone> CharStream<'src>
    for TextCharStream<'src, I>
{
    fn peek(&self) -> Option<char> {
        self.peeked.map(|c| c.1)
    }

    fn step(&mut self) -> bool {
        let c = self.peek().expect("step called on EOF");
        self.current_pos.advance(c);
        match self.current_chars.next() {
            Some(c) => {
                self.peeked = Some(c);
                false
            }
            None => {
                self.peeked = self.next_new();
                true
            }
        }
    }

    fn char_idx(&self) -> usize {
        match self.peeked {
            Some((ofs, _)) => ofs,
            None => self.current_chunk.text.len(),
        }
    }

    fn pos(&self) -> Pos {
        Pos {
            origin: self.current_chunk.pos.origin,
            text_pos: self.current_pos.pos(self.peek()),
        }
    }

    fn current_chunk(&self) -> &'src TextChunk {
        self.current_chunk
    }
}

pub struct ChunkState<'chr, 'src, I: CharStream<'src>> {
    chars: &'chr mut I,
    _pd: PhantomData<&'src ()>,

    start_idx: usize,
    start_pos: Pos,
}

impl<'chr, 'src, I: CharStream<'src>> ChunkState<'chr, 'src, I> {
    fn new(chars: &'chr mut I) -> Self {
        let start_idx = chars.char_idx();
        let start_pos = chars.pos();
        Self {
            chars,
            _pd: PhantomData,
            start_idx,
            start_pos,
        }
    }

    pub fn peek(&mut self) -> Option<char> {
        self.chars.peek()
    }

    fn get_subchunk_from_str(&self, chunk: &TextChunk, end: usize) -> Option<TextChunk> {
        let range = self.start_idx..end;

        if range.start == range.end {
            return None;
        }

        Some(TextChunk {
            pos: self.start_pos,
            text: chunk.text.substr(range),
        })
    }

    fn get_subchunk(&self, end: usize) -> Option<TextChunk> {
        self.get_subchunk_from_str(self.chars.current_chunk(), end)
    }

    fn reset_subchunk(&mut self) {
        self.start_idx = self.chars.char_idx();
        self.start_pos = self.chars.pos();
    }

    /// Advance one char, including it in the next subchunk.
    /// Return Some(subchunk) on chunk boundary.
    #[must_use]
    pub fn advance(&mut self) -> Option<TextChunk> {
        self.chars.peek().expect("advance called on EOF");
        let chunk = self.chars.current_chunk();

        match self.chars.step() {
            false => None,
            true => {
                let chunk = self.get_subchunk_from_str(chunk, chunk.text.len());
                self.reset_subchunk();
                chunk
            }
        }
    }

    #[must_use]
    /// Advance one char, ommiting it from the next subchunk.
    /// Return pending subchunk.
    pub fn skip(&mut self) -> Option<TextChunk> {
        let chunk = self.get_subchunk(self.chars.char_idx());
        self.chars.next();
        self.reset_subchunk();
        chunk
    }

    #[must_use]
    /// Consume stream, return pending subchunk
    fn flush(&mut self) -> Option<TextChunk> {
        let chunk = self.get_subchunk(self.chars.char_idx());
        self.reset_subchunk();
        chunk
    }

    pub fn writer<F: FnMut(TextChunk)>(self, f: F) -> FnChunkWriter<'chr, 'src, I, F> {
        FnChunkWriter {
            sub: self,
            writer: f,
        }
    }

    pub fn pos(&self) -> Pos {
        self.chars.pos()
    }
}

// To hide away all those lifetimes
pub trait ChunkWriter<'src> {
    type Chars: CharStream<'src> + ?Sized;

    fn peek(&mut self) -> Option<char>;

    /// Advance one char, including it in the next subchunk.
    fn advance(&mut self);

    /// Advance one char, ommiting it from the next subchunk.
    fn skip(&mut self);

    /// Consume stream, return pending subchunk
    fn flush(&mut self);

    fn pos(&self) -> Pos;

    fn chars(&self) -> &Self::Chars;

    fn skipped(&mut self) -> SubChunkSkip<'_, 'src, Self> {
        SubChunkSkip {
            sub: self,
            _pd: PhantomData,
        }
    }

    fn check_any(&mut self) {
        self.peek().expect("consume_any called on empty chunk");
    }

    fn check(&mut self, expected: char) {
        assert_eq!(self.peek(), Some(expected))
    }

    fn advance_while(&mut self, mut pred: impl FnMut(char) -> bool) -> usize {
        let mut cnt = 0;
        while let Some(c) = self.peek() {
            if !pred(c) {
                break;
            }
            cnt += 1;
            self.advance();
        }
        cnt
    }

    fn skip_while(&mut self, mut pred: impl FnMut(char) -> bool) -> usize {
        let mut cnt = 0;
        while let Some(c) = self.peek() {
            if !pred(c) {
                break;
            }
            cnt += 1;
            self.skip();
        }
        cnt
    }
}

pub struct FnChunkWriter<'chr, 'src, I: CharStream<'src>, F: FnMut(TextChunk)> {
    sub: ChunkState<'chr, 'src, I>,
    writer: F,
}

impl<'src, I: CharStream<'src>, F: FnMut(TextChunk)> FnChunkWriter<'_, 'src, I, F> {
    fn writer(&mut self) -> impl FnOnce(TextChunk) + '_ {
        |chunk| (self.writer)(chunk)
    }
}

impl<'src, I: CharStream<'src>, F: FnMut(TextChunk)> ChunkWriter<'src>
    for FnChunkWriter<'_, 'src, I, F>
{
    type Chars = I;

    fn peek(&mut self) -> Option<char> {
        self.sub.peek()
    }

    /// Advance one char, including it in the next subchunk.
    fn advance(&mut self) {
        self.sub.advance().then(self.writer());
    }

    /// Advance one char, ommiting it from the next subchunk.
    fn skip(&mut self) {
        self.sub.skip().then(self.writer())
    }

    /// Consume stream, return pending subchunk
    fn flush(&mut self) {
        self.sub.flush().then(self.writer())
    }

    fn pos(&self) -> Pos {
        self.sub.pos()
    }

    fn chars(&self) -> &Self::Chars {
        self.sub.chars
    }
}

pub struct SubChunkSkip<'a, 'src, T: ChunkWriter<'src> + ?Sized> {
    sub: &'a mut T,
    _pd: PhantomData<&'src ()>,
}

impl<'src, T: ChunkWriter<'src>> ChunkWriter<'src> for SubChunkSkip<'_, 'src, T> {
    type Chars = T::Chars;

    fn peek(&mut self) -> Option<char> {
        self.sub.peek()
    }

    fn advance(&mut self) {
        self.sub.skip();
    }

    fn skip(&mut self) {
        self.sub.skip();
    }

    fn pos(&self) -> Pos {
        self.sub.pos()
    }

    // No need to do anything
    fn flush(&mut self) {}

    fn chars(&self) -> &Self::Chars {
        self.sub.chars()
    }
}