Spans from custom tokens.

[aDotInTheVoid]

I've been trying to write a winnow parser that uses a custom lexer.
I've been having some trouble having the parser correctly pick up the spans that the lexer
assigns to tokens.

NB: A full, compilling version of this can be found here

Part 0: The lexer.

The lexer exposes a relativly simple interface.

enum TokenKind {
    LParen,
    RParen,
    // actually more
}

pub struct Token {
    kind: TokenKind,
    span: ops::Range<usize>,
}

fn lex(src: &str) -> Vec<Token> {
    ...
}

Each token carries it's own span, which is needed to preserve information as
whitespace is stripped, as I don't want to need to skip whitespace everywhere in
the parser. (This is one of the motivations for using a custom lexer).

#[test]
fn test_lexer() {
    let toks = lex("( () () )");
    //              0123456789
    let exp = expect![[r#"
        [
            LParen(0..1),
            LParen(2..3),
            RParen(3..4),
            LParen(5..6),
            RParen(6..7),
            RParen(8..9),
        ]
    "#]];
    exp.assert_debug_eq(&toks);
}

In order to use this with the primitives, I've used the following implementations.

// A TokenKind is a parser that accepts any token with that kind.
impl<'a, 'b, I, E> Parser<I, Token, E> for TokenKind
where
    I: Stream<Token = Token> + StreamIsPartial,
    E: ParserError<I>,
{
    fn parse_next(&mut self, input: &mut I) -> PResult<Token, E> {
        one_of(*self).parse_next(input)
    }
}
// Helper for the above implementation.
impl ContainsToken<Token> for TokenKind {
    fn contains_token(&self, token: Token) -> bool {
        token.kind == *self
    }
}

To demonstrate this, I'll be using a simple sexpr parser. It's generic over the
stream that it parsers, as long as it implements the standard traits.

struct SExpr {
    inner: Vec<SExpr>,
    span: ops::Range<usize>,
}

fn sexpr<I>(i: &mut I) -> PResult<SExpr>
where
    I: Stream<Token = Token> + Location,
    TokenKind: Parser<I, Token, ContextError>,
{
    delimited(LParen, repeat(0.., sexpr), RParen)
        .with_span()
        .map(|(inner, span)| SExpr { inner, span })
        .parse_next(i)
}

I'll be using the same input in all examples. Ideally it should parse to the tree below.

( () () )
0123456789

(Sexp 0..9
    (Sexp 2..4)
    (Sexp 5..7))

Approach 1: Using Located.

My first guess was to just use the Located type.

type Input1<'a> = Located<&'a [Token]>;

Unfortunatly this doesn't give the results I want.

#[test]
fn test_input_1() {
    let toks = lex("( () () )");
    let input: Input1 = Located::new(&toks);

    let tree = sexpr::<Input1>.parse(input).unwrap();
    let exp = expect![[r#"
        SExpr {
            inner: [
                SExpr {
                    inner: [],
                    span: 1..3,
                },
                SExpr {
                    inner: [],
                    span: 3..5,
                },
            ],
            span: 0..6,
        }
    "#]];
    exp.assert_debug_eq(&tree);
}

The spans here are token offsets, not byte offsets (as provided from the lexer).

Approach 2: Manually implementing Location.

Given that Located only advances the .location() by 1 each time a token is consumed, it's clear I need a different Location impl. (unless I want to implement .offset, but that seems like a bad idea, given that finding the offset between any 2 tokens is harder that finding the location of a single token).

#[derive(Debug, Clone)]
struct Input2<'a> {
    inner: &'a [Token],
}

impl Location for Input2<'_> {
    fn location(&self) -> usize {
        match self.inner.first() {
            Some(t) => t.span.start,
            None => panic!("Can't find location of empty input"),
        }
    }
}

This also requires a load of boilerplate to implement Offset, Stream and StreamIsPartial (not shown here).

#[test]
fn test_input_2() {
    let toks = lex("( () () )");

    let input = Input2 { inner: &toks };

    sexpr.parse(input).unwrap();
}

Unfortunatly it doesn't work, as it can't find the location for the end of input, as there's no token left.

---- test_input_2 stdout ----
thread 'test_input_2' panicked at mcve.rs:144:21:
Can't find location of empty input
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

Approach 3: State!

But given that we've got our own struct, we can just store it there, and access it later.

#[derive(Debug, Clone)]
struct Input3<'a> {
    inner: &'a [Token],
    end_byte_offset: usize,
}

impl Location for Input3<'_> {
    fn location(&self) -> usize {
        match self.inner.first() {
            Some(t) => t.span.start,
            None => self.end_byte_offset,
        }
    }
}

This also needs a Offset, Stream and StreamIsPartial implementation.

This works alot better.

#[test]
fn test_input_3() {
    let src = "( () () )";
    let toks = lex(src);

    let input = Input3 {
        inner: &toks,
        end_byte_offset: toks.last().map_or(|t| t.span.end, 0),
    };

    let tree = sexpr.parse(input).unwrap();

    let exp = expect![[r#"
        SExpr {
            inner: [
                SExpr {
                    inner: [],
                    span: 2..5,
                },
                SExpr {
                    inner: [],
                    span: 5..8,
                },
            ],
            span: 0..9,
        }
    "#]];
    exp.assert_debug_eq(&tree);
}

unfortunatly it's got the spans wrong. Because we only reference the .start of spans, it can't see that the lexer has decided that the whitespace isn't a part of the token.

More pictorally it parsers into

( () () )
  ^^^
  |  ^^^
  |  |
^^^^^^^^^

instead of

( () () )
  ^^
  |  ^^
  |  |
^^^^^^^^^

Conclusion.

At this point, I investigated what was going on with with_span, to see if it could be cajoled into doing what I wanted.
it's impl:

    fn parse_next(&mut self, input: &mut I) -> PResult<(O, Range<usize>), E> {
        let start = input.location();
        self.parser.parse_next(input).map(move |output| {
            let end = input.location();
            (output, (start..end))
        })
    }

The problem here is that it uses the same call for the start and end.

If the parser is in this situation:

     end    start
     v      v
aaa bb      cccc ddd
            ^
            parser is here

we want it to use the lower location for the .end of bb and the upper for
the .start of cccc.

Doing this would require changing the Location trait, and the implmentation of
.span and .with_span. While in theory I could do this myself, I thaught it'd
be worth asking your thaught on this, and how much of this makes sense to live
upstream.

[epage]

Yes, Location was very much designed with Lexing (+Parsing) in mind. It was designed to be a better nom_locate. I would be open to explorations of what the design could be for Parsing of lexed tokens. It also seems like less of a gap to fulfill to get spans when working with lexed tokens that carry their own span. I also suspect the design considerations are a bit different than lexing.

[aDotInTheVoid]

Here's what I ended up going with: I gave up on using the Location trait entirely, and took advantage of using a custom input type.

pub struct Input<'toks, 'acx> {
    source: &'toks [lexer::Token],
    offset: usize,
    
    pub(crate) acx: AstCtxt<'acx>, // not relevant here, used to store arena
}

impl<'toks, 'acx> Input<'toks, 'acx> {
    /// ```text
    ///     prev_token_high─┐   ┌─next_token_low
    ///                     ▼   ▼
    ///   a   bb   ccc   dddd   eeee    | source
    /// 0   1    2     3      4         | offset
    /// ```
    fn prev_token_high(&self) -> usize {
        self.source[self.offset - 1].span().hi()
    }

    fn next_token_low(&self) -> usize {
        match self.source.get(self.offset) {
            Some(tok) => tok.span().lo(),
            None => {
                //                            ▼ offset
                // a   bb   ccc   dddd   eeee
                debug_assert_eq!(self.offset, self.source.len());
                self.source.last().unwrap().span().hi()
            }
        }
    }
}

pub(crate) fn with_span<'acx, 'tok, O, P, E>(
    mut parser: P,
) -> impl Parser<Input<'tok, 'acx>, (O, Span), E>
where
    P: Parser<Input<'tok, 'acx>, O, E>,
{
    move |i: &mut Input<'tok, 'acx>| {
        let start = i.next_token_low();
        parser.parse_next(i).map(|output| {
            let end = i.prev_token_high();
            (output, Span::new(start, end))
        })
    }
}

Input here also has a implementation of Stream that needs to take into account to only show tokens after offset, (unless it get's reset to a checkpoint). Interestingly many of those Stream method's never end up getting called (outside of tests), but I implement them anyway as I don't want to rely on this. It might be possible to move them to a different trait though, but I've not looked into this.

[epage]

In #591 (reply in thread), I'm considering expanding Location to offer similar information as you put on your custom Input.

[epage]

btw starting a higher level conversation on improving the lexing+parsing story over at #591