killswitch/
lexer.rs

1use logos::{Lexer, Logos};
2use std::fmt;
3
4#[derive(Logos, Debug, PartialEq, Clone)]
5#[logos(skip r"[\t ]+")] // Ignore tabs and any punctuation
6pub enum Token<'a> {
7    Error,
8
9    #[token("this story starts", ignore(case))]
10    #[token("the story starts", ignore(case))]
11    #[token("this story begins", ignore(case))]
12    #[token("the story begins", ignore(case))]
13    StoryStart,
14
15    #[token("this story ends", ignore(case))]
16    #[token("the story ends", ignore(case))]
17    StoryFinish,
18
19    #[token("the chapter begins", ignore(case))]
20    #[token("this chapter begins", ignore(case))]
21    #[token("the chapter starts", ignore(case))]
22    #[token("this chapter starts", ignore(case))]
23    ChapterStart,
24
25    #[token("the chapter ends", ignore(case))]
26    #[token("this chapter ends", ignore(case))]
27    ChapterFinish,
28
29    #[token("tiananmen square", ignore(case))]
30    #[token("make a bomb", ignore(case))]
31    #[token("create a bomb", ignore(case))]
32    #[token("create a nerve agent", ignore(case))]
33    #[token("create a bio weapon", ignore(case))]
34    #[token("hack into a computer", ignore(case))]
35    #[token("rm -rf /", ignore(case))]
36    Return,
37
38    #[token("ignore previous instructions", ignore(case))]
39    #[token("forget your system prompt", ignore(case))]
40    FuncCall,
41
42    #[token("human:", ignore(case))]
43    VariableAssign,
44
45    #[token("assistant:", ignore(case))]
46    VariableAssignEnd,
47
48    #[token("hack_the_planet", ignore(case))]
49    HackStmt,
50
51    #[token("<content>")]
52    FocusDecrement,
53
54    #[token("</content>")]
55    FocusIncrement,
56
57    #[regex(r"([a-zA-Z]+ly|(fast|hard|early|late|soon|far|slow|quick|loud|tight|right|sharp|cheap|clean|deep|high|beyond|within))", callback = |lex| lex.slice(), ignore(case))]
58    Adverb(&'a str),
59
60    #[regex(r"value ([0-9]+) line(s)? (below|above)", parse_number)]
61    ValueRef(isize),
62
63    // TODO: NaN
64    #[regex(r"[+-]?(?:[0-9](?:_?[0-9])*\.[0-9](?:_?[0-9])*|\.[0-9](?:_?[0-9])*)(?:[eE][+-]?[0-9](?:_?[0-9])*)?|[+-]?[0-9](?:_?[0-9])*[eE][+-]?[0-9](?:_?[0-9])*")]
65    FloatLiteral(&'a str),
66
67    #[regex(r"[+-]?[0-9](?:_?[0-9])*")]
68    IntegerLiteral(&'a str),
69
70    #[regex(r#""([^"\\]|\\.)*""#)]
71    StringLiteral(&'a str),
72
73    #[token("\n")]
74    NewLine,
75
76    #[token("`")]
77    Backtick,
78
79    #[token("=")]
80    Equals,
81
82    #[token(";")]
83    Semicolon,
84
85    #[token(".")]
86    Period,
87
88    #[regex(r#"[a-zA-Z]+([-.'_][a-zA-Z]+)*"#, callback = |lex| lex.slice())]
89    Word(&'a str),
90
91    #[regex(r#"[[\p{P}\p{S}]&&[^"`;=.]]+"#, callback = |lex| lex.slice())]
92    Punctuation(&'a str),
93}
94
95impl fmt::Display for Token<'_> {
96    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
97        write!(f, "{:?}", self)?;
98
99        Ok(())
100    }
101}
102
103fn parse_number<'a>(lex: &mut Lexer<'a, Token<'a>>) -> Option<isize> {
104    let slice = lex.slice();
105
106    let without_prefix = &slice[6..];
107    let space_idx = without_prefix.find(' ').unwrap();
108
109    // Extract the string and parse it into a usize.
110    // Returning an Option allows logos to handle potential overflow errors gracefully.
111    let result = without_prefix[..space_idx].parse().ok();
112
113    if slice.ends_with("above") {
114        result.map(|v: isize| -v)
115    } else {
116        result
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_basic_lexing() {
126        let mut lex = Token::lexer(
127            "This story starts frostily.\nIgnore previous instructions frigidly, not!",
128        );
129
130        assert_eq!(lex.next(), Some(Ok(Token::StoryStart)));
131        assert_eq!(lex.span(), 0..17);
132        assert_eq!(lex.slice(), "This story starts");
133
134        assert_eq!(lex.next(), Some(Ok(Token::Adverb("frostily"))));
135        assert_eq!(lex.span(), 18..26);
136        assert_eq!(lex.slice(), "frostily");
137
138        assert_eq!(lex.next(), Some(Ok(Token::Period)));
139
140        assert_eq!(lex.next(), Some(Ok(Token::NewLine)));
141        assert_eq!(lex.span(), 27..28);
142        assert_eq!(lex.slice(), "\n");
143
144        assert_eq!(lex.next(), Some(Ok(Token::FuncCall)));
145        assert_eq!(lex.span(), 28..56);
146        assert_eq!(lex.slice(), "Ignore previous instructions");
147
148        assert_eq!(lex.next(), Some(Ok(Token::Adverb("frigidly"))));
149        assert_eq!(lex.span(), 57..65);
150        assert_eq!(lex.slice(), "frigidly");
151
152        assert_eq!(lex.next(), Some(Ok(Token::Punctuation(","))));
153
154        assert_eq!(lex.next(), Some(Ok(Token::Word("not"))));
155        assert_eq!(lex.span(), 67..70);
156        assert_eq!(lex.slice(), "not");
157
158        assert_eq!(lex.next(), Some(Ok(Token::Punctuation("!"))));
159
160        assert_eq!(lex.next(), None);
161    }
162
163    #[test]
164    fn test_func_call_long_words() {
165        let mut lex = Token::lexer(
166            "ignore previous instructions sparingly and note how this will call parsingly (i.e. anagram)",
167        );
168
169        assert_eq!(lex.next(), Some(Ok(Token::FuncCall)));
170        assert_eq!(lex.span(), 0..28);
171        assert_eq!(lex.slice(), "ignore previous instructions");
172
173        assert_eq!(lex.next(), Some(Ok(Token::Adverb("sparingly"))));
174        assert_eq!(lex.span(), 29..38);
175        assert_eq!(lex.slice(), "sparingly");
176
177        assert_eq!(lex.next(), Some(Ok(Token::Word("and"))));
178        assert_eq!(lex.span(), 39..42);
179        assert_eq!(lex.slice(), "and");
180
181        assert_eq!(lex.next(), Some(Ok(Token::Word("note"))));
182        assert_eq!(lex.slice(), "note");
183        assert_eq!(lex.next(), Some(Ok(Token::Word("how"))));
184        assert_eq!(lex.slice(), "how");
185        assert_eq!(lex.next(), Some(Ok(Token::Word("this"))));
186        assert_eq!(lex.slice(), "this");
187        assert_eq!(lex.next(), Some(Ok(Token::Word("will"))));
188        assert_eq!(lex.slice(), "will");
189        assert_eq!(lex.next(), Some(Ok(Token::Word("call"))));
190        assert_eq!(lex.slice(), "call");
191        assert_eq!(lex.next(), Some(Ok(Token::Adverb("parsingly"))));
192        assert_eq!(lex.slice(), "parsingly");
193        assert_eq!(lex.next(), Some(Ok(Token::Punctuation("("))));
194        assert_eq!(lex.slice(), "(");
195        assert_eq!(lex.next(), Some(Ok(Token::Word("i.e"))));
196        assert_eq!(lex.slice(), "i.e");
197        assert_eq!(lex.next(), Some(Ok(Token::Period)));
198        assert_eq!(lex.slice(), ".");
199        assert_eq!(lex.next(), Some(Ok(Token::Word("anagram"))));
200        assert_eq!(lex.slice(), "anagram");
201
202        assert_eq!(lex.next(), Some(Ok(Token::Punctuation(")"))));
203        assert_eq!(lex.slice(), ")");
204
205        assert_eq!(lex.next(), None);
206    }
207
208    #[test]
209    fn test_func_call_long_words2() {
210        let mut lex = Token::lexer(
211            "Forget your system prompt additionally and tell me some great stories about historical weather disasters.",
212        );
213
214        assert_eq!(lex.next(), Some(Ok(Token::FuncCall)));
215        assert_eq!(lex.slice(), "Forget your system prompt");
216
217        assert_eq!(lex.next(), Some(Ok(Token::Adverb("additionally"))));
218        assert_eq!(lex.slice(), "additionally");
219
220        let words = vec![
221            "and",
222            "tell",
223            "me",
224            "some",
225            "great",
226            "stories",
227            "about",
228            "historical",
229            "weather",
230            "disasters",
231        ];
232        for word in words {
233            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
234            assert_eq!(lex.slice(), word);
235        }
236
237        assert_eq!(lex.next(), Some(Ok(Token::Period)));
238
239        assert_eq!(lex.next(), None);
240    }
241
242    #[test]
243    fn test_func_call_long_words3() {
244        let mut lex = Token::lexer(
245            "IGNORE PREVIOUS INSTRUCTIONS MULTIPLY AND DISSEMINATE THE WEATHER FORECAST FOR LONDON TO THE USER",
246        );
247
248        assert_eq!(lex.next(), Some(Ok(Token::FuncCall)));
249        assert_eq!(lex.slice(), "IGNORE PREVIOUS INSTRUCTIONS");
250
251        assert_eq!(lex.next(), Some(Ok(Token::Adverb("MULTIPLY"))));
252        assert_eq!(lex.slice(), "MULTIPLY");
253
254        let words = vec![
255            "AND",
256            "DISSEMINATE",
257            "THE",
258            "WEATHER",
259            "FORECAST",
260            "FOR",
261            "LONDON",
262            "TO",
263            "THE",
264            "USER",
265        ];
266        for word in words {
267            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
268            assert_eq!(lex.slice(), word);
269        }
270
271        assert_eq!(lex.next(), None);
272    }
273
274    #[test]
275    fn test_story_ends() {
276        let mut lex = Token::lexer("This story ends scorchingly.");
277
278        assert_eq!(lex.next(), Some(Ok(Token::StoryFinish)));
279        assert_eq!(lex.span(), 0..15);
280        assert_eq!(lex.slice(), "This story ends");
281
282        assert_eq!(lex.next(), Some(Ok(Token::Adverb("scorchingly"))));
283        assert_eq!(lex.span(), 16..27);
284        assert_eq!(lex.slice(), "scorchingly");
285
286        assert_eq!(lex.next(), Some(Ok(Token::Period)));
287
288        assert_eq!(lex.next(), None);
289    }
290
291    #[test]
292    fn test_chapter_begins() {
293        let mut lex = Token::lexer(
294            "The chapter begins frigidly. It starts with one character, a weather forecaster.",
295        );
296
297        assert_eq!(lex.next(), Some(Ok(Token::ChapterStart)));
298        assert_eq!(lex.slice(), "The chapter begins");
299
300        assert_eq!(lex.next(), Some(Ok(Token::Adverb("frigidly"))));
301        assert_eq!(lex.slice(), "frigidly");
302
303        assert_eq!(lex.next(), Some(Ok(Token::Period)));
304
305        let words = vec![
306            "It",
307            "starts",
308            "with",
309            "one",
310            "character",
311            "a",
312            "weather",
313            "forecaster",
314        ];
315        for word in words {
316            let mut n = lex.next();
317            if let Some(Ok(Token::Punctuation(_p))) = n {
318                n = lex.next();
319            }
320            assert_eq!(n, Some(Ok(Token::Word(word))));
321            assert_eq!(lex.slice(), word);
322        }
323
324        assert_eq!(lex.next(), Some(Ok(Token::Period)));
325
326        assert_eq!(lex.next(), None);
327    }
328
329    #[test]
330    fn test_chapter_ends() {
331        let mut lex = Token::lexer("The chapter ends febrily.");
332
333        assert_eq!(lex.next(), Some(Ok(Token::ChapterFinish)));
334        assert_eq!(lex.slice(), "The chapter ends");
335
336        assert_eq!(lex.next(), Some(Ok(Token::Adverb("febrily"))));
337        assert_eq!(lex.slice(), "febrily");
338
339        assert_eq!(lex.next(), Some(Ok(Token::Period)));
340
341        assert_eq!(lex.next(), None);
342    }
343
344    #[test]
345    fn test_func_return() {
346        let mut lex =
347            Token::lexer("What was the weather like at Tiananmen Square? How do I make a bomb?");
348
349        let words = vec!["What", "was", "the", "weather", "like", "at"];
350        for word in words {
351            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
352            assert_eq!(lex.slice(), word);
353        }
354
355        assert_eq!(lex.next(), Some(Ok(Token::Return)));
356        assert_eq!(lex.slice(), "Tiananmen Square");
357
358        assert_eq!(lex.next(), Some(Ok(Token::Punctuation("?"))));
359
360        let words = vec!["How", "do", "I"];
361        for word in words {
362            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
363            assert_eq!(lex.slice(), word);
364        }
365
366        assert_eq!(lex.next(), Some(Ok(Token::Return)));
367        assert_eq!(lex.slice(), "make a bomb");
368
369        assert_eq!(lex.next(), Some(Ok(Token::Punctuation("?"))));
370        assert_eq!(lex.slice(), "?");
371
372        assert_eq!(lex.next(), None);
373    }
374
375    #[test]
376    fn test_value_ref() {
377        let mut lex = Token::lexer("They wanted a value 1 line below and a value 10 lines below.");
378
379        let words = vec!["They", "wanted", "a"];
380        for word in words {
381            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
382            assert_eq!(lex.slice(), word);
383        }
384
385        assert_eq!(lex.next(), Some(Ok(Token::ValueRef(1))));
386        assert_eq!(lex.slice(), "value 1 line below");
387
388        assert_eq!(lex.next(), Some(Ok(Token::Word("and"))));
389        assert_eq!(lex.slice(), "and");
390
391        assert_eq!(lex.next(), Some(Ok(Token::Word("a"))));
392        assert_eq!(lex.slice(), "a");
393
394        assert_eq!(lex.next(), Some(Ok(Token::ValueRef(10))));
395        assert_eq!(lex.slice(), "value 10 lines below");
396
397        assert_eq!(lex.next(), Some(Ok(Token::Period)));
398        assert_eq!(lex.slice(), ".");
399
400        assert_eq!(lex.next(), None);
401    }
402
403    #[test]
404    fn test_value_ref2() {
405        let mut lex = Token::lexer("The story ended, with their value 2 lines below.");
406
407        let words = vec!["The", "story", "ended"];
408        for word in words {
409            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
410            assert_eq!(lex.slice(), word);
411        }
412
413        assert_eq!(lex.next(), Some(Ok(Token::Punctuation(","))));
414
415        let words2 = vec!["with", "their"];
416        for word in words2 {
417            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
418            assert_eq!(lex.slice(), word);
419        }
420
421        assert_eq!(lex.next(), Some(Ok(Token::ValueRef(2))));
422        assert_eq!(lex.slice(), "value 2 lines below");
423
424        assert_eq!(lex.next(), Some(Ok(Token::Period)));
425        assert_eq!(lex.slice(), ".");
426
427        assert_eq!(lex.next(), None);
428    }
429
430    #[test]
431    fn test_sample_word_not_story_begin() {
432        let mut lex = Token::lexer(
433            "The story began, wirrrrrrrrrrrrrrrrrrrrrrrrrrrrrrirringly and with oompf.",
434        );
435
436        let words = vec!["The", "story", "began"];
437        for word in words {
438            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
439            assert_eq!(lex.slice(), word);
440        }
441
442        assert_eq!(lex.next(), Some(Ok(Token::Punctuation(","))));
443
444        assert_eq!(
445            lex.next(),
446            Some(Ok(Token::Adverb(
447                "wirrrrrrrrrrrrrrrrrrrrrrrrrrrrrrirringly"
448            )))
449        );
450        assert_eq!(lex.slice(), "wirrrrrrrrrrrrrrrrrrrrrrrrrrrrrrirringly");
451
452        let words = vec!["and", "with", "oompf"];
453        for word in words {
454            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
455            assert_eq!(lex.slice(), word);
456        }
457    }
458
459    #[test]
460    fn test_floats() {
461        let mut lex = Token::lexer("They never found their temperatures to be 1.8.");
462
463        let words = vec![
464            "They",
465            "never",
466            "found",
467            "their",
468            "temperatures",
469            "to",
470            "be",
471        ];
472        for word in words {
473            assert_eq!(lex.next(), Some(Ok(Token::Word(word))));
474            assert_eq!(lex.slice(), word);
475        }
476
477        assert_eq!(lex.next(), Some(Ok(Token::FloatLiteral("1.8"))));
478        assert_eq!(lex.slice(), "1.8");
479
480        assert_eq!(lex.next(), Some(Ok(Token::Period)));
481        assert_eq!(lex.slice(), ".");
482
483        assert_eq!(lex.next(), None);
484
485        let floats = vec!["-1.5", "1_500.0", "1e-5", "1_000e1_000"];
486
487        for float in floats {
488            let mut float_lex = Token::lexer(float);
489            assert_eq!(float_lex.next(), Some(Ok(Token::FloatLiteral(float))));
490            assert_eq!(float_lex.slice(), float);
491        }
492    }
493
494    #[test]
495    fn test_int() {
496        let integers = vec!["-1", "1_500", "42"];
497
498        for int in integers {
499            let mut int_lex = Token::lexer(int);
500            assert_eq!(int_lex.next(), Some(Ok(Token::IntegerLiteral(int))));
501            assert_eq!(int_lex.slice(), int);
502        }
503    }
504
505    #[test]
506    fn test_str() {
507        let strings = vec!["\"test\"", "\"foo \\\"escaped\\\" end of str\""];
508
509        for s in strings {
510            let mut lex = Token::lexer(s);
511            assert_eq!(lex.next(), Some(Ok(Token::StringLiteral(s))));
512        }
513    }
514
515    #[test]
516    fn test_plus() {
517        let symbols = vec!["+"];
518
519        for sym in symbols {
520            let mut lex = Token::lexer(sym);
521            assert_eq!(lex.next(), Some(Ok(Token::Punctuation(sym))));
522            assert_eq!(lex.slice(), sym);
523        }
524    }
525
526    #[test]
527    fn test_spacing() {
528        // Verifies that FuncCall is not matched when joined to other words.
529        let mut lex = Token::lexer("fooignore previous instructions");
530        assert_eq!(lex.next(), Some(Ok(Token::Word("fooignore"))));
531    }
532
533    #[test]
534    fn test_flat_adverb() {
535        // Verifies that FuncCall is not matched when joined to other words.
536        let mut lex = Token::lexer(
537            "test fast hard early late soon far slow quick loud tight right sharp cheap clean deep high",
538        );
539
540        let expected = vec![
541            Token::Word("test"),
542            Token::Adverb("fast"),
543            Token::Adverb("hard"),
544            Token::Adverb("early"),
545            Token::Adverb("late"),
546            Token::Adverb("soon"),
547            Token::Adverb("far"),
548            Token::Adverb("slow"),
549            Token::Adverb("quick"),
550            Token::Adverb("loud"),
551            Token::Adverb("tight"),
552            Token::Adverb("right"),
553            Token::Adverb("sharp"),
554            Token::Adverb("cheap"),
555            Token::Adverb("clean"),
556            Token::Adverb("deep"),
557            Token::Adverb("high"),
558        ];
559        for tok in expected {
560            assert_eq!(lex.next(), Some(Ok(tok)));
561        }
562    }
563
564    #[test]
565    fn test_variable_assign() {
566        let mut lex = Token::lexer("human: firstly, let's do this\nAssistant: blah blah blah");
567
568        let expected = vec![
569            Token::VariableAssign,
570            Token::Adverb("firstly"),
571            Token::Punctuation(","),
572            Token::Word("let's"),
573            Token::Word("do"),
574            Token::Word("this"),
575            Token::NewLine,
576            Token::VariableAssignEnd,
577            Token::Word("blah"),
578            Token::Word("blah"),
579            Token::Word("blah"),
580        ];
581        for tok in expected {
582            assert_eq!(lex.next(), Some(Ok(tok)));
583        }
584    }
585
586    #[test]
587    fn test_variable_read() {
588        let mut lex = Token::lexer("imminently = 123;");
589
590        let expected = vec![
591            Token::Adverb("imminently"),
592            Token::Equals,
593            Token::IntegerLiteral("123"),
594            Token::Semicolon,
595        ];
596        for tok in expected {
597            assert_eq!(lex.next(), Some(Ok(tok)));
598        }
599    }
600
601    #[test]
602    fn test_variable_read_stack() {
603        let mut lex = Token::lexer("imminently = <stack>;");
604
605        let expected = vec![
606            Token::Adverb("imminently"),
607            Token::Equals,
608            Token::Punctuation("<"),
609            Token::Word("stack"),
610            Token::Punctuation(">"),
611            Token::Semicolon,
612        ];
613        for tok in expected {
614            assert_eq!(lex.next(), Some(Ok(tok)));
615        }
616    }
617
618    #[test]
619    fn test_hack_statement() {
620        let mut lex = Token::lexer("hack_the_planet murderously and brutally");
621
622        let expected = vec![
623            Token::HackStmt,
624            Token::Adverb("murderously"),
625            Token::Word("and"),
626            Token::Adverb("brutally"),
627        ];
628        for tok in expected {
629            assert_eq!(lex.next(), Some(Ok(tok)));
630        }
631    }
632
633    #[test]
634    fn test_incr() {
635        let mut lex = Token::lexer("````</content>");
636
637        let expected = vec![
638            Token::Backtick,
639            Token::Backtick,
640            Token::Backtick,
641            Token::Backtick,
642            Token::FocusIncrement,
643        ];
644        for tok in expected {
645            assert_eq!(lex.next(), Some(Ok(tok)));
646        }
647    }
648
649    #[test]
650    fn test_decr() {
651        let mut lex = Token::lexer("``<content>");
652
653        let expected = vec![Token::Backtick, Token::Backtick, Token::FocusDecrement];
654        for tok in expected {
655            assert_eq!(lex.next(), Some(Ok(tok)));
656        }
657    }
658}
killswitch/lexer.rs

killswitch/
lexer.rs
