killswitch/
parser.rs

1use chumsky::{
2    input::{Stream, ValueInput},
3    prelude::*,
4    span::SimpleSpan,
5};
6use logos::Logos;
7use std::fmt;
8
9use crate::lexer::Token;
10
11pub type Span = std::ops::Range<usize>;
12pub type Spanned<N> = (N, Span);
13
14#[derive(Clone, Debug, PartialEq)]
15pub enum Node {
16    Story(String, Vec<Spanned<Self>>),
17    Chapter(String, Vec<Spanned<Self>>),
18    FuncCall(String, Vec<Spanned<Self>>), // ident, nodes captured so they can be counted and used for collect_values
19    IntLiteral(i64),
20    FloatLiteral(f64),
21    StringLiteral(String),
22    FuncReturn,
23    ValueRef(isize), // count of "lines below"
24    Adverb(String),  // can be referenced by ValueRef
25    Word(String),
26    VariableAssign(String),
27    VariableRead(String, Box<Spanned<Self>>),
28    Hack(String, Option<String>),
29    FocusChange(isize),
30}
31
32impl fmt::Display for Node {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            Node::Story(ident, nodes) => {
36                write!(f, "Story({}, ", ident)?;
37                for (i, (node, _span)) in nodes.iter().enumerate() {
38                    if i != 0 {
39                        write!(f, ", ")?;
40                    }
41                    write!(f, "{}", node)?;
42                }
43                write!(f, ")")
44            }
45            Node::Chapter(ident, nodes) => {
46                write!(f, "Chapter({}, ", ident)?;
47                for (i, (node, _span)) in nodes.iter().enumerate() {
48                    if i != 0 {
49                        write!(f, ", ")?;
50                    }
51                    write!(f, "{}", node)?;
52                }
53                write!(f, ")")
54            }
55            Node::FuncCall(ident, words) => write!(f, "FuncCall({}, {})", ident, words.len()),
56            Node::IntLiteral(val) => write!(f, "{}", val),
57            Node::FloatLiteral(val) => write!(f, "{}", val),
58            Node::StringLiteral(val) => write!(f, "{}", val),
59            Node::FuncReturn => write!(f, "FuncReturn"),
60            Node::ValueRef(lines) => write!(f, "ValueRef({})", lines),
61            Node::Adverb(adverb) => write!(f, "Adverb({})", adverb),
62            Node::Word(w) => write!(f, "Word({})", w),
63            Node::VariableAssign(adverb) => write!(f, "VariableAssign({})", adverb),
64            Node::VariableRead(ident, children) => {
65                write!(f, "VariableRead({}, {:?})", ident, children)
66            }
67            Node::Hack(a, b) => write!(f, "Hack({}, {:?})", a, b),
68            Node::FocusChange(t) => write!(f, "FocusChange({})", t),
69        }
70    }
71}
72
73// S can be SimpleSpan<usize> or Span
74//
75// Fun fact: I caused a Rust ICE here https://share.gemini.google/rVzs8vGPTE73
76impl Node {
77    pub fn spanned<S: Into<Span>>(self, span: S) -> (Self, Span) {
78        (self, span.into())
79    }
80}
81
82pub trait Unspanned<N> {
83    fn unspanned(&self) -> &N;
84}
85
86// Just a little helper to convert a Spanned node to
87// Node. We need the trait because we cannot define
88// functions on a tuple directly.
89impl Unspanned<Node> for Spanned<Node> {
90    fn unspanned(&self) -> &Node {
91        &self.0
92    }
93}
94
95// Based on example in https://github.com/zesterer/chumsky/blob/main/examples/logos.rs#L73.
96pub fn parser<'tok, 'src: 'tok, I>()
97-> impl Parser<'tok, I, Vec<Spanned<Node>>, extra::Err<Rich<'tok, Token<'src>>>>
98where
99    I: ValueInput<'tok, Token = Token<'src>, Span = SimpleSpan>,
100{
101    let adverb_to_ident = select! {
102        Token::Adverb(ident) => ident,
103    };
104
105    let word_like = any().filter(|t| {
106        matches!(
107            t,
108            Token::Word(_)
109                | Token::Adverb(_)
110                | Token::IntegerLiteral(_)
111                | Token::FloatLiteral(_)
112                | Token::StringLiteral(_)
113        )
114    });
115
116    let ignored = any().filter(|t| {
117        matches!(
118            t,
119            Token::NewLine
120                | Token::Punctuation(_)
121                | Token::Equals
122                | Token::Semicolon
123                | Token::Period
124        )
125    });
126
127    let single_node = recursive(|value| {
128        let word_like_atom = select! {
129            Token::Word(w) = e => Node::Word(w.to_owned()).spanned(e.span()),
130            Token::Adverb(w) = e => Node::Adverb(w.to_owned()).spanned(e.span()),
131            Token::FloatLiteral(s) = e => Node::FloatLiteral(s.parse().unwrap()).spanned(e.span()),
132            Token::IntegerLiteral(s) = e => Node::IntLiteral(s.parse().unwrap()).spanned(e.span()),
133            Token::StringLiteral(s) = e => Node::StringLiteral(s.to_owned()).spanned(e.span()),
134        };
135
136        // Ref: https://docs.rs/chumsky/latest/chumsky/macro.select.html
137        let atom = select! {
138            Token::Word(w) = e => Node::Word(w.to_owned()).spanned(e.span()),
139            Token::Return = e => Node::FuncReturn.spanned(e.span()),
140            Token::ValueRef(line_count) = e => Node::ValueRef(line_count).spanned(e.span()),
141        }
142        .or(word_like_atom);
143
144        // Idea: perhaps we allow the adverb used to not be an anagram, if the word
145        // count rules are satisfied, but if not we do allow anagrams (but the whole sentence
146        // can't have any adverbs then)
147        //
148        // We currently use a somewhat complex rule: pick the first adverb in the words list,
149        // if there is no adverb amongst the words then use the first word. This handles
150        // anagrams put at the start of the function call and allows adverbs to be called
151        // from other positions.
152        let func_call = just(Token::FuncCall)
153            .then(
154                word_like_atom
155                    .repeated()
156                    .collect::<Vec<_>>()
157                    .separated_by(any().filter(|t| matches!(t, Token::Punctuation(_))))
158                    .collect::<Vec<Vec<_>>>()
159                    .map(|words| words.into_iter().flatten().collect::<Vec<_>>()),
160            )
161            .try_map_with(|(_, mut words), e| {
162                let first_adverb_idx = words
163                    .clone()
164                    .into_iter()
165                    .position(|t| matches!(t.unspanned(), Node::Adverb(_)));
166                // TODO: Show warning when multiple adverbs exist?
167                let first_adverb = first_adverb_idx.map(|idx| words.remove(idx));
168
169                match (first_adverb, words.get(0)) {
170                    (Some((Node::Adverb(ident), _)), _) => {
171                        Ok(Node::FuncCall(ident.to_owned(), words).spanned(e.span()))
172                    }
173                    (_, Some((Node::Word(ident), _))) => {
174                        Ok(Node::FuncCall(ident.to_owned(), words[1..].to_vec()).spanned(e.span()))
175                    }
176                    _ => Err(Rich::custom(
177                        e.span(),
178                        "Function call requires an identifier in the form of a word or adverb",
179                    )),
180                }
181            });
182
183        let variable_assignment = just(Token::VariableAssign)
184            .then(
185                word_like_atom
186                    .repeated()
187                    .collect::<Vec<_>>()
188                    .separated_by(ignored)
189                    .collect::<Vec<Vec<_>>>()
190                    .map(|lines| lines.into_iter().flatten().collect::<Vec<_>>()),
191            )
192            .then_ignore(just(Token::VariableAssignEnd))
193            .try_map_with(|(_, words), e| {
194                let first_adverb = words
195                    .into_iter()
196                    .find(|t| matches!(t.unspanned(), Node::Adverb(_)));
197
198                // TODO: Store the `words` in the VariableAssign as children, in case
199                // there are any value refs there.
200                match first_adverb {
201                    Some((Node::Adverb(ident), _)) => {
202                        Ok(Node::VariableAssign(ident.to_owned()).spanned(e.span()))
203                    }
204                    _ => Err(Rich::custom(
205                        e.span(),
206                        "Variable assignment requires at least one adverb",
207                    )),
208                }
209            });
210
211        let any_punctuation = select! { Token::Punctuation(_) => () };
212        let surrounded_atom = word_like_atom
213            .clone()
214            .delimited_by(any_punctuation, any_punctuation);
215
216        let variable_read = adverb_to_ident
217            .then_ignore(just(Token::Equals))
218            .then(surrounded_atom.or(word_like_atom))
219            .then_ignore(just(Token::Semicolon))
220            .map_with(|(ident, child), e| {
221                Node::VariableRead(ident.to_owned(), Box::new(child)).spanned(e.span())
222            });
223
224        let hack_stmt = just(Token::HackStmt)
225            .then(word_like.repeated().collect::<Vec<_>>())
226            .try_map_with(|(_, words), e| {
227                let mut adverbs = words.into_iter().filter(|t| matches!(t, Token::Adverb(_)));
228                match (adverbs.next(), adverbs.next()) {
229                    (Some(Token::Adverb(a)), _b_token @ Some(Token::Adverb(b))) => {
230                        Ok(Node::Hack(a.to_string(), Some(b.to_string())).spanned(e.span()))
231                    }
232                    (Some(Token::Adverb(a)), _) => {
233                        Ok(Node::Hack(a.to_string(), None).spanned(e.span()))
234                    }
235                    _ => Err(Rich::custom(
236                        e.span(),
237                        "Hack statement requires at least one adverb",
238                    )),
239                }
240            });
241
242        let focus_decr = just(Token::Backtick)
243            .repeated()
244            .at_least(1)
245            .collect::<Vec<_>>()
246            .then(just(Token::FocusDecrement))
247            .map_with(|(t, _), e| Node::FocusChange(t.len() as isize).spanned(e.span()));
248
249        let focus_incr = just(Token::Backtick)
250            .repeated()
251            .at_least(1)
252            .collect::<Vec<_>>()
253            .then(just(Token::FocusIncrement))
254            .map_with(|(t, _), e| Node::FocusChange(-(t.len() as isize)).spanned(e.span()));
255
256        let backtick =
257            just(Token::Backtick).map_with(|_, e| Node::Word("`".to_owned()).spanned(e.span()));
258
259        let node_list = value
260            .clone()
261            .separated_by(ignored.repeated())
262            .allow_leading()
263            .allow_trailing()
264            .collect::<Vec<_>>();
265
266        let story = adverb_to_ident
267            .then(node_list.clone())
268            .delimited_by(just(Token::StoryStart), just(Token::StoryFinish))
269            .then_ignore(adverb_to_ident)
270            .map_with(|(ident, stmts), e| Node::Story(ident.to_owned(), stmts).spanned(e.span()));
271
272        let chapter = adverb_to_ident
273            .then(node_list.clone())
274            .delimited_by(just(Token::ChapterStart), just(Token::ChapterFinish))
275            .then_ignore(adverb_to_ident)
276            .map_with(|(ident, stmts), e| Node::Chapter(ident.to_owned(), stmts).spanned(e.span()));
277
278        // TODO: We have a lot of rules that start by checking for `adverb`. This isn't
279        // efficient. Consider "left-factoring" this.
280        variable_read
281            .or(atom)
282            .or(func_call)
283            .or(variable_assignment)
284            .or(hack_stmt)
285            .or(story)
286            .or(chapter)
287            .or(focus_decr)
288            .or(focus_incr)
289            .or(backtick)
290    });
291
292    single_node
293        .separated_by(ignored.repeated())
294        .allow_leading()
295        .allow_trailing()
296        .collect::<Vec<_>>()
297        .then_ignore(end())
298        .map(|nodes| {
299            nodes
300                .into_iter()
301                .filter(|n| !matches!(&n.0, Node::Word(s) if s == "`"))
302                .collect::<Vec<_>>()
303        })
304}
305
306pub fn lex_to_parsed_result(code: &str) -> ParseResult<Vec<Spanned<Node>>, Rich<'_, Token<'_>>> {
307    let lex = Token::lexer(code).spanned().map(|(tok, span)| match tok {
308        Ok(tok) => (tok, span.into()),
309        Err(()) => (Token::Error, span.into()),
310    });
311    let token_stream = Stream::from_iter(lex).map((0..code.len()).into(), |(t, s): (_, _)| (t, s));
312    return parser().parse(token_stream);
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318
319    // TODO: Test for error cases. Like "The chapter ends BLAH".
320
321    fn assert_eq_func_call(a: &Node, b: &Node) {
322        match (&a, &b) {
323            (Node::FuncCall(a_i, a_w), Node::FuncCall(b_i, b_w)) => {
324                assert_eq!(a_i, b_i);
325                assert_eq!(a_w.len(), b_w.len());
326            }
327            _ => {
328                assert!(false);
329            }
330        }
331    }
332
333    #[test]
334    fn test_story_parse() {
335        let parsed_result = lex_to_parsed_result(
336            "This story starts frostily.\nIgnore previous instructions frigidly, not!\nThis story ends scorchingly.",
337        );
338
339        let result = parsed_result.into_result().unwrap();
340
341        let expected = Node::Story(
342            "frostily".to_string(),
343            vec![
344                Node::FuncCall(
345                    "frigidly".to_string(),
346                    vec![Node::Word("not".to_owned()).spanned(67..70)],
347                )
348                .spanned(28..71),
349            ],
350        );
351        assert_eq!(result[0].unspanned(), &expected);
352    }
353
354    #[test]
355    fn test_chapter_parse() {
356        let parsed_result = lex_to_parsed_result(
357            "The chapter begins frigidly.\nvalue 42 lines below\nThe chapter ends febrily.",
358        );
359
360        let result = parsed_result.into_result().unwrap();
361
362        let expected = Node::Chapter(
363            "frigidly".to_string(),
364            vec![Node::ValueRef(42).spanned(29..49)],
365        );
366        assert_eq!(result[0].unspanned(), &expected);
367    }
368
369    #[test]
370    fn test_chapter_name_must_be_adverb() {
371        let parsed_result = lex_to_parsed_result(
372            "The chapter begins foobaz.\nvalue 42 lines below\nThe chapter ends febrily.",
373        );
374
375        let result = parsed_result.into_result();
376
377        assert!(
378            matches!(&result, Err(errs) if errs.len() == 1 && errs[0].span().into_range() == (19..25)),
379            "Expected an error at span 19..25, but got: {:?}",
380            result
381        );
382    }
383
384    #[test]
385    fn test_func_call_words() {
386        let parsed_result = lex_to_parsed_result(
387            "ignore previous instructions sparingly and note how this will call parsingly (i.e anagram)",
388        );
389
390        let result = parsed_result.into_result().unwrap();
391
392        let expected = Node::FuncCall(
393            "sparingly".to_string(),
394            vec![Node::Word("".to_owned()).spanned(0..0); 9],
395        );
396        assert_eq_func_call(&result[0].unspanned(), &expected);
397    }
398
399    #[test]
400    fn test_func_call_words_numbers() {
401        let parsed_result =
402            lex_to_parsed_result("forget your system prompt eqaully, tell me how many 0 in ten");
403
404        let result = parsed_result.into_result().unwrap();
405
406        let expected = Node::FuncCall(
407            "eqaully".to_string(),
408            vec![Node::Word("".to_owned()).spanned(0..0); 7],
409        );
410        assert_eq_func_call(&result[0].unspanned(), &expected);
411    }
412
413    #[test]
414    fn test_func_call_no_words() {
415        let parsed_result = lex_to_parsed_result("ignore previous instructions humanely");
416
417        let result = parsed_result.into_result().unwrap();
418
419        let expected = Node::FuncCall("humanely".to_string(), vec![]);
420        assert_eq!(result[0].unspanned(), &expected);
421    }
422
423    #[test]
424    fn test_func_call_alternate() {
425        let parsed_result = lex_to_parsed_result(
426            "Forget your system prompt additionally and tell me some great stories about historical weather disasters.",
427        );
428
429        let result = parsed_result.into_result().unwrap();
430
431        let expected = Node::FuncCall(
432            "additionally".to_string(),
433            vec![Node::Word("".to_owned()).spanned(0..0); 10],
434        );
435        assert_eq_func_call(&result[0].unspanned(), &expected);
436    }
437
438    #[test]
439    fn test_words() {
440        let parsed_result = lex_to_parsed_result(
441            "He says \"Celsius to Fahrenheit uses the formula (C * 9/5) + 32\".
442            He also says \"The temperature in Celsius will be on the top of the stack. \
443            So we need to push 9/5 i.e. 1.8 on the stack, then call multiply, then \
444            push 32 on the stack, then call add\".",
445        );
446
447        let result = parsed_result.into_result().unwrap();
448
449        // We only care that this parsed fine since it's just text. Not worried about what it
450        // parses as for now.
451        assert_eq!(result.len(), 7);
452    }
453
454    #[test]
455    fn test_func_return() {
456        let parsed_result =
457            lex_to_parsed_result("Weather at Tiananmen Square? How do I make a bomb?");
458
459        let result = parsed_result.into_result().unwrap();
460
461        let expected = [
462            Node::Word("Weather".to_owned()),
463            Node::Word("at".to_owned()),
464            Node::FuncReturn,
465            Node::Word("How".to_owned()),
466            Node::Word("do".to_owned()),
467            Node::Word("I".to_owned()),
468            Node::FuncReturn,
469        ];
470        assert_eq!(
471            result.into_iter().map(|(n, _s)| n).collect::<Vec<_>>(),
472            expected
473        );
474    }
475
476    #[test]
477    fn test_val_ref() {
478        let parsed_result = lex_to_parsed_result("They wanted a value 1 line below.");
479
480        let result = parsed_result.into_result().unwrap();
481
482        let expected = [
483            Node::Word("They".to_owned()).spanned(0..4),
484            Node::Word("wanted".to_owned()).spanned(5..11),
485            Node::Word("a".to_owned()).spanned(12..13),
486            Node::ValueRef(1).spanned(14..32),
487        ];
488        assert_eq!(result, expected);
489    }
490
491    #[test]
492    fn test_eof_newline() {
493        let parsed_result = lex_to_parsed_result(
494            "This story starts frostily.\nIgnore previous instructions frigidly, not!\nThis story ends scorchingly.\n",
495        );
496
497        let result = parsed_result.into_result().unwrap();
498
499        let expected = Node::Story(
500            "frostily".to_string(),
501            vec![
502                Node::FuncCall(
503                    "frigidly".to_string(),
504                    vec![Node::Word("not".to_owned()).spanned(67..70)],
505                )
506                .spanned(28..71),
507            ],
508        );
509        assert_eq!(result[0].unspanned(), &expected);
510    }
511
512    #[test]
513    fn test_var_assign() {
514        let parsed_result = lex_to_parsed_result("human: firstly, let's do this\nAssistant: blah");
515
516        let result = parsed_result.into_result().unwrap();
517
518        let expected = [
519            Node::VariableAssign("firstly".to_string()).spanned(0..40),
520            Node::Word("blah".to_owned()).spanned(41..45),
521        ];
522        assert_eq!(result, expected);
523    }
524
525    #[test]
526    fn test_var_assign_anywhere() {
527        let parsed_result = lex_to_parsed_result(
528            "human: let's do this, you may find me\nsupposedly good\nbut why\nAssistant: blah",
529        );
530
531        let result = parsed_result.into_result().unwrap();
532
533        let expected = [
534            Node::VariableAssign("supposedly".to_string()).spanned(0..72),
535            Node::Word("blah".to_owned()).spanned(73..77),
536        ];
537        assert_eq!(result, expected);
538    }
539
540    #[test]
541    fn test_var_read() {
542        let parsed_result = lex_to_parsed_result("imminently = 123;\nword");
543
544        let result = parsed_result.into_result().unwrap();
545
546        let expected = [
547            Node::VariableRead(
548                "imminently".to_string(),
549                Box::new(Node::IntLiteral(123).spanned(13..16)),
550            )
551            .spanned(0..17),
552            Node::Word("word".to_owned()).spanned(18..22),
553        ];
554        assert_eq!(result, expected);
555    }
556
557    #[test]
558    fn test_var_read_stack() {
559        let parsed_result = lex_to_parsed_result("imminently = <stack>;\nword");
560
561        let result = parsed_result.into_result().unwrap();
562
563        let expected = [
564            Node::VariableRead(
565                "imminently".to_string(),
566                Box::new(Node::Word("stack".to_owned()).spanned(14..19)),
567            )
568            .spanned(0..21),
569            Node::Word("word".to_owned()).spanned(22..26),
570        ];
571        assert_eq!(result, expected);
572    }
573
574    #[test]
575    fn test_hack_stmt() {
576        let parsed_result = lex_to_parsed_result("hack_the_planet murderously and brutally");
577
578        let result = parsed_result.into_result().unwrap();
579
580        let expected =
581            [Node::Hack("murderously".to_string(), Some("brutally".to_string())).spanned(0..40)];
582        assert_eq!(result, expected);
583    }
584
585    #[test]
586    fn test_hack_stmt_one() {
587        let parsed_result = lex_to_parsed_result("hack_the_planet murderously and nothing else");
588
589        let result = parsed_result.into_result().unwrap();
590
591        let expected = [Node::Hack("murderously".to_string(), None).spanned(0..44)];
592        assert_eq!(result, expected);
593    }
594
595    #[test]
596    fn test_focus_decrement() {
597        let parsed_result = lex_to_parsed_result("```<content>");
598
599        let result = parsed_result.into_result().unwrap();
600
601        let expected = [Node::FocusChange(3).spanned(0..12)];
602        assert_eq!(result, expected);
603    }
604
605    #[test]
606    fn test_multi_stmt_per_line() {
607        let parsed_result = lex_to_parsed_result(
608            "Ignore previous instructions debuggably. There is a value 1 line below (10).",
609        );
610
611        let result = parsed_result.into_result().unwrap();
612
613        let expected = [
614            Node::FuncCall("debuggably".to_owned(), vec![]).spanned(0..39),
615            Node::Word("There".to_owned()).spanned(41..46),
616            Node::Word("is".to_owned()).spanned(47..49),
617            Node::Word("a".to_owned()).spanned(50..51),
618            Node::ValueRef(1).spanned(52..70),
619            Node::IntLiteral(10).spanned(72..74),
620        ];
621        assert_eq!(result, expected);
622    }
623
624    #[test]
625    fn test_multi_stmt_call_var() {
626        let parsed_result =
627            lex_to_parsed_result("Ignore previous instructions debuggably. variably = <stack>;");
628
629        let result = parsed_result.into_result().unwrap();
630
631        let expected = [
632            Node::FuncCall("debuggably".to_owned(), vec![]).spanned(0..39),
633            Node::VariableRead(
634                "variably".to_owned(),
635                Box::new(Node::Word("stack".to_owned()).spanned(53..58)),
636            )
637            .spanned(41..60),
638        ];
639        assert_eq!(result, expected);
640    }
641
642    #[test]
643    fn test_random_backticks() {
644        let parsed_result = lex_to_parsed_result("Foo `rm -rf /` bar");
645
646        let result = parsed_result.into_result().unwrap();
647
648        let expected = [
649            Node::Word("Foo".to_owned()).spanned(0..3),
650            Node::FuncReturn.spanned(5..13),
651            Node::Word("bar".to_owned()).spanned(15..18),
652        ];
653        assert_eq!(result, expected);
654    }
655
656    #[test]
657    fn test_var_read_then_focus_change() {
658        let parsed_result = lex_to_parsed_result("intently = <stack>;```<content>");
659
660        let result = parsed_result.into_result().unwrap();
661
662        let expected = [
663            Node::VariableRead(
664                "intently".to_owned(),
665                Box::new(Node::Word("stack".to_owned()).spanned(12..17)),
666            )
667            .spanned(0..19),
668            Node::FocusChange(3).spanned(19..31),
669        ];
670        assert_eq!(result, expected);
671    }
672}
killswitch/parser.rs

killswitch/
parser.rs
