trustformers_tokenizers/
code_tokenizer.rs

1use serde::{Deserialize, Serialize};
2use std::collections::{HashMap, HashSet};
3use trustformers_core::errors::Result;
4use trustformers_core::traits::{TokenizedInput, Tokenizer};
5
6/// Supported programming languages
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub enum Language {
9    Rust,
10    Python,
11    JavaScript,
12    TypeScript,
13    Java,
14    CSharp,
15    CPlusPlus,
16    C,
17    Go,
18    Ruby,
19    PHP,
20    Swift,
21    Kotlin,
22    Scala,
23    Haskell,
24    Clojure,
25    SQL,
26    HTML,
27    CSS,
28    JSON,
29    XML,
30    YAML,
31    Markdown,
32    Shell,
33    PowerShell,
34    R,
35    Matlab,
36}
37
38impl Language {
39    /// Get file extensions for the language
40    pub fn extensions(&self) -> &'static [&'static str] {
41        match self {
42            Language::Rust => &["rs"],
43            Language::Python => &["py", "pyx", "pyi", "pyw"],
44            Language::JavaScript => &["js", "jsx", "mjs", "cjs"],
45            Language::TypeScript => &["ts", "tsx", "d.ts"],
46            Language::Java => &["java"],
47            Language::CSharp => &["cs"],
48            Language::CPlusPlus => &["cpp", "cxx", "cc", "hpp", "hxx", "hh"],
49            Language::C => &["c", "h"],
50            Language::Go => &["go"],
51            Language::Ruby => &["rb", "rbx", "rjs", "gemspec"],
52            Language::PHP => &["php", "phtml", "php3", "php4", "php5"],
53            Language::Swift => &["swift"],
54            Language::Kotlin => &["kt", "kts"],
55            Language::Scala => &["scala", "sc"],
56            Language::Haskell => &["hs", "lhs"],
57            Language::Clojure => &["clj", "cljs", "cljc", "edn"],
58            Language::SQL => &["sql"],
59            Language::HTML => &["html", "htm", "xhtml"],
60            Language::CSS => &["css", "scss", "sass", "less"],
61            Language::JSON => &["json", "jsonl", "ndjson"],
62            Language::XML => &["xml", "xsd", "xsl", "xslt"],
63            Language::YAML => &["yaml", "yml"],
64            Language::Markdown => &["md", "markdown", "mdown", "mkd"],
65            Language::Shell => &["sh", "bash", "zsh", "fish"],
66            Language::PowerShell => &["ps1", "psm1", "psd1"],
67            Language::R => &["r", "R"],
68            Language::Matlab => &["m"],
69        }
70    }
71
72    /// Get keywords for the language
73    pub fn keywords(&self) -> &'static [&'static str] {
74        match self {
75            Language::Rust => &[
76                "as", "break", "const", "continue", "crate", "else", "enum", "extern", "false",
77                "fn", "for", "if", "impl", "in", "let", "loop", "match", "mod", "move", "mut",
78                "pub", "ref", "return", "self", "Self", "static", "struct", "super", "trait",
79                "true", "type", "unsafe", "use", "where", "while", "async", "await", "dyn",
80            ],
81            Language::Python => &[
82                "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class",
83                "continue", "def", "del", "elif", "else", "except", "finally", "for", "from",
84                "global", "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
85                "raise", "return", "try", "while", "with", "yield",
86            ],
87            Language::JavaScript | Language::TypeScript => &[
88                "break",
89                "case",
90                "catch",
91                "class",
92                "const",
93                "continue",
94                "debugger",
95                "default",
96                "delete",
97                "do",
98                "else",
99                "export",
100                "extends",
101                "false",
102                "finally",
103                "for",
104                "function",
105                "if",
106                "import",
107                "in",
108                "instanceof",
109                "new",
110                "null",
111                "return",
112                "super",
113                "switch",
114                "this",
115                "throw",
116                "true",
117                "try",
118                "typeof",
119                "var",
120                "void",
121                "while",
122                "with",
123                "yield",
124                "let",
125                "static",
126                "enum",
127                "implements",
128                "package",
129                "protected",
130                "interface",
131                "private",
132                "public",
133                "async",
134                "await",
135            ],
136            Language::Java => &[
137                "abstract",
138                "assert",
139                "boolean",
140                "break",
141                "byte",
142                "case",
143                "catch",
144                "char",
145                "class",
146                "const",
147                "continue",
148                "default",
149                "do",
150                "double",
151                "else",
152                "enum",
153                "extends",
154                "final",
155                "finally",
156                "float",
157                "for",
158                "goto",
159                "if",
160                "implements",
161                "import",
162                "instanceof",
163                "int",
164                "interface",
165                "long",
166                "native",
167                "new",
168                "package",
169                "private",
170                "protected",
171                "public",
172                "return",
173                "short",
174                "static",
175                "strictfp",
176                "super",
177                "switch",
178                "synchronized",
179                "this",
180                "throw",
181                "throws",
182                "transient",
183                "try",
184                "void",
185                "volatile",
186                "while",
187            ],
188            Language::CSharp => &[
189                "abstract",
190                "as",
191                "base",
192                "bool",
193                "break",
194                "byte",
195                "case",
196                "catch",
197                "char",
198                "checked",
199                "class",
200                "const",
201                "continue",
202                "decimal",
203                "default",
204                "delegate",
205                "do",
206                "double",
207                "else",
208                "enum",
209                "event",
210                "explicit",
211                "extern",
212                "false",
213                "finally",
214                "fixed",
215                "float",
216                "for",
217                "foreach",
218                "goto",
219                "if",
220                "implicit",
221                "in",
222                "int",
223                "interface",
224                "internal",
225                "is",
226                "lock",
227                "long",
228                "namespace",
229                "new",
230                "null",
231                "object",
232                "operator",
233                "out",
234                "override",
235                "params",
236                "private",
237                "protected",
238                "public",
239                "readonly",
240                "ref",
241                "return",
242                "sbyte",
243                "sealed",
244                "short",
245                "sizeof",
246                "stackalloc",
247                "static",
248                "string",
249                "struct",
250                "switch",
251                "this",
252                "throw",
253                "true",
254                "try",
255                "typeof",
256                "uint",
257                "ulong",
258                "unchecked",
259                "unsafe",
260                "ushort",
261                "using",
262                "virtual",
263                "void",
264                "volatile",
265                "while",
266            ],
267            Language::Go => &[
268                "break",
269                "case",
270                "chan",
271                "const",
272                "continue",
273                "default",
274                "defer",
275                "else",
276                "fallthrough",
277                "for",
278                "func",
279                "go",
280                "goto",
281                "if",
282                "import",
283                "interface",
284                "map",
285                "package",
286                "range",
287                "return",
288                "select",
289                "struct",
290                "switch",
291                "type",
292                "var",
293            ],
294            _ => &[], // Add more languages as needed
295        }
296    }
297
298    /// Get comment patterns for the language
299    pub fn comment_patterns(&self) -> CommentPatterns {
300        match self {
301            Language::Rust
302            | Language::JavaScript
303            | Language::TypeScript
304            | Language::Java
305            | Language::CSharp
306            | Language::CPlusPlus
307            | Language::Go
308            | Language::Swift
309            | Language::Kotlin
310            | Language::Scala => CommentPatterns {
311                line_comment: Some("//"),
312                block_comment: Some(("/*", "*/")),
313                doc_comment: Some("///"),
314            },
315            Language::Python | Language::Ruby | Language::Shell => CommentPatterns {
316                line_comment: Some("#"),
317                block_comment: None,
318                doc_comment: Some("#"),
319            },
320            Language::C => CommentPatterns {
321                line_comment: None,
322                block_comment: Some(("/*", "*/")),
323                doc_comment: None,
324            },
325            Language::HTML | Language::XML => CommentPatterns {
326                line_comment: None,
327                block_comment: Some(("<!--", "-->")),
328                doc_comment: None,
329            },
330            Language::CSS => CommentPatterns {
331                line_comment: None,
332                block_comment: Some(("/*", "*/")),
333                doc_comment: None,
334            },
335            Language::SQL => CommentPatterns {
336                line_comment: Some("--"),
337                block_comment: Some(("/*", "*/")),
338                doc_comment: None,
339            },
340            Language::Haskell => CommentPatterns {
341                line_comment: Some("--"),
342                block_comment: Some(("{-", "-}")),
343                doc_comment: Some("-- |"),
344            },
345            _ => CommentPatterns {
346                line_comment: None,
347                block_comment: None,
348                doc_comment: None,
349            },
350        }
351    }
352
353    /// Detect language from file extension
354    pub fn from_extension(ext: &str) -> Option<Language> {
355        let ext = ext.to_lowercase();
356        [
357            Language::Rust,
358            Language::Python,
359            Language::JavaScript,
360            Language::TypeScript,
361            Language::Java,
362            Language::CSharp,
363            Language::CPlusPlus,
364            Language::C,
365            Language::Go,
366            Language::Ruby,
367            Language::PHP,
368            Language::Swift,
369            Language::Kotlin,
370            Language::Scala,
371            Language::Haskell,
372            Language::Clojure,
373            Language::SQL,
374            Language::HTML,
375            Language::CSS,
376            Language::JSON,
377            Language::XML,
378            Language::YAML,
379            Language::Markdown,
380            Language::Shell,
381            Language::PowerShell,
382            Language::R,
383            Language::Matlab,
384        ]
385        .into_iter()
386        .find(|&lang| lang.extensions().contains(&ext.as_str()))
387    }
388}
389
390/// Comment patterns for a language
391#[derive(Debug, Clone)]
392pub struct CommentPatterns {
393    pub line_comment: Option<&'static str>,
394    pub block_comment: Option<(&'static str, &'static str)>,
395    pub doc_comment: Option<&'static str>,
396}
397
398/// Token types for code
399#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
400pub enum CodeTokenType {
401    Keyword,
402    Identifier,
403    Literal(LiteralType),
404    Operator,
405    Punctuation,
406    Comment,
407    Whitespace,
408    String,
409    Number,
410    Unknown,
411}
412
413/// Types of literals
414#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
415pub enum LiteralType {
416    String,
417    Character,
418    Integer,
419    Float,
420    Boolean,
421    Null,
422}
423
424/// A code token with type information
425#[derive(Debug, Clone, Serialize, Deserialize)]
426pub struct CodeToken {
427    pub text: String,
428    pub token_type: CodeTokenType,
429    pub position: TokenPosition,
430    pub language: Language,
431}
432
433/// Position information for a token
434#[derive(Debug, Clone, Serialize, Deserialize)]
435pub struct TokenPosition {
436    pub line: usize,
437    pub column: usize,
438    pub start_offset: usize,
439    pub end_offset: usize,
440}
441
442/// Configuration for code tokenization
443#[derive(Debug, Clone, Serialize, Deserialize)]
444pub struct CodeTokenizerConfig {
445    pub language: Option<Language>,
446    pub preserve_whitespace: bool,
447    pub preserve_comments: bool,
448    pub include_position_info: bool,
449    pub normalize_identifiers: bool,
450    pub max_token_length: Option<usize>,
451    pub custom_keywords: Option<HashSet<String>>,
452}
453
454impl Default for CodeTokenizerConfig {
455    fn default() -> Self {
456        Self {
457            language: None,
458            preserve_whitespace: false,
459            preserve_comments: true,
460            include_position_info: false,
461            normalize_identifiers: false,
462            max_token_length: Some(128),
463            custom_keywords: None,
464        }
465    }
466}
467
468/// Code tokenizer implementation
469pub struct CodeTokenizer {
470    config: CodeTokenizerConfig,
471    keywords: HashSet<String>,
472    token_to_id: HashMap<String, u32>,
473    id_to_token: HashMap<u32, String>,
474    special_tokens: HashMap<String, u32>,
475}
476
477impl CodeTokenizer {
478    /// Create a new code tokenizer
479    pub fn new(config: CodeTokenizerConfig) -> Self {
480        let mut tokenizer = Self {
481            config,
482            keywords: HashSet::new(),
483            token_to_id: HashMap::new(),
484            id_to_token: HashMap::new(),
485            special_tokens: HashMap::new(),
486        };
487
488        tokenizer.initialize_vocabulary();
489        tokenizer
490    }
491
492    /// Create tokenizer for a specific language
493    pub fn for_language(language: Language) -> Self {
494        let config = CodeTokenizerConfig {
495            language: Some(language),
496            ..Default::default()
497        };
498        Self::new(config)
499    }
500
501    /// Initialize vocabulary with common tokens
502    fn initialize_vocabulary(&mut self) {
503        let mut next_id = 0u32;
504
505        // Add special tokens
506        for special in &[
507            "[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[BOS]", "[EOS]",
508        ] {
509            self.add_token(special, &mut next_id);
510            self.special_tokens.insert(special.to_string(), next_id - 1);
511        }
512
513        // Add language keywords
514        if let Some(language) = self.config.language {
515            for keyword in language.keywords() {
516                self.keywords.insert(keyword.to_string());
517                self.add_token(keyword, &mut next_id);
518            }
519        }
520
521        // Add custom keywords
522        if let Some(custom_keywords) = &self.config.custom_keywords {
523            let keywords_to_add: Vec<String> = custom_keywords.iter().cloned().collect();
524            for keyword in keywords_to_add {
525                self.keywords.insert(keyword.clone());
526                self.add_token(&keyword, &mut next_id);
527            }
528        }
529
530        // Add common operators and punctuation
531        for op in &[
532            "+", "-", "*", "/", "%", "=", "==", "!=", "<", ">", "<=", ">=", "&&", "||", "!", "&",
533            "|", "^", "~", "<<", ">>", "++", "--", "+=", "-=", "*=", "/=", "%=", "(", ")", "[",
534            "]", "{", "}", ";", ",", ".", ":", "::", "->", "=>", "?",
535        ] {
536            self.add_token(op, &mut next_id);
537        }
538
539        // Add common literals
540        for literal in &["true", "false", "null", "undefined", "nil", "None"] {
541            self.add_token(literal, &mut next_id);
542        }
543    }
544
545    /// Add a token to the vocabulary
546    fn add_token(&mut self, token: &str, next_id: &mut u32) {
547        if !self.token_to_id.contains_key(token) {
548            self.token_to_id.insert(token.to_string(), *next_id);
549            self.id_to_token.insert(*next_id, token.to_string());
550            *next_id += 1;
551        }
552    }
553
554    /// Tokenize code into structured tokens
555    pub fn tokenize_code(&self, code: &str) -> Result<Vec<CodeToken>> {
556        let language = self.config.language.unwrap_or(Language::JavaScript);
557        let comment_patterns = language.comment_patterns();
558
559        let mut tokens = Vec::new();
560        let mut current_line = 1;
561        let mut current_column = 1;
562        let mut char_indices = code.char_indices().peekable();
563
564        while let Some((start_offset, ch)) = char_indices.next() {
565            let token_start_line = current_line;
566            let token_start_column = current_column;
567
568            // Update position
569            if ch == '\n' {
570                current_line += 1;
571                current_column = 1;
572            } else {
573                current_column += 1;
574            }
575
576            // Skip whitespace (unless preserving)
577            if ch.is_whitespace() {
578                if self.config.preserve_whitespace {
579                    let (text, end_offset) =
580                        self.consume_whitespace(&mut char_indices, start_offset, ch);
581                    tokens.push(CodeToken {
582                        text,
583                        token_type: CodeTokenType::Whitespace,
584                        position: TokenPosition {
585                            line: token_start_line,
586                            column: token_start_column,
587                            start_offset,
588                            end_offset,
589                        },
590                        language,
591                    });
592                }
593                continue;
594            }
595
596            // Handle comments
597            if let Some(token) = self.try_parse_comment(
598                &mut char_indices,
599                start_offset,
600                ch,
601                &comment_patterns,
602                token_start_line,
603                token_start_column,
604                language,
605            )? {
606                if self.config.preserve_comments {
607                    tokens.push(token);
608                }
609                continue;
610            }
611
612            // Handle string literals
613            if ch == '"'
614                || ch == '\''
615                || (ch == '`' && matches!(language, Language::JavaScript | Language::TypeScript))
616            {
617                let token = self.parse_string_literal(
618                    &mut char_indices,
619                    start_offset,
620                    ch,
621                    token_start_line,
622                    token_start_column,
623                    language,
624                )?;
625                tokens.push(token);
626                continue;
627            }
628
629            // Handle numeric literals
630            if ch.is_ascii_digit()
631                || (ch == '.'
632                    && char_indices.peek().map(|(_, c)| c.is_ascii_digit()).unwrap_or(false))
633            {
634                let token = self.parse_numeric_literal(
635                    &mut char_indices,
636                    start_offset,
637                    ch,
638                    token_start_line,
639                    token_start_column,
640                    language,
641                )?;
642                tokens.push(token);
643                continue;
644            }
645
646            // Handle identifiers and keywords
647            if ch.is_alphabetic() || ch == '_' || ch == '$' {
648                let token = self.parse_identifier(
649                    &mut char_indices,
650                    start_offset,
651                    ch,
652                    token_start_line,
653                    token_start_column,
654                    language,
655                )?;
656                tokens.push(token);
657                continue;
658            }
659
660            // Handle operators and punctuation
661            let token = self.parse_operator_or_punctuation(
662                &mut char_indices,
663                start_offset,
664                ch,
665                token_start_line,
666                token_start_column,
667                language,
668            )?;
669            tokens.push(token);
670        }
671
672        Ok(tokens)
673    }
674
675    /// Consume whitespace characters
676    fn consume_whitespace(
677        &self,
678        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
679        start_offset: usize,
680        first_char: char,
681    ) -> (String, usize) {
682        let mut text = String::new();
683        text.push(first_char);
684        let mut end_offset = start_offset;
685
686        while let Some((offset, ch)) = char_indices.peek() {
687            if ch.is_whitespace() {
688                text.push(*ch);
689                end_offset = *offset;
690                char_indices.next();
691            } else {
692                break;
693            }
694        }
695
696        (text, end_offset)
697    }
698
699    /// Try to parse a comment
700    fn try_parse_comment(
701        &self,
702        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
703        start_offset: usize,
704        first_char: char,
705        patterns: &CommentPatterns,
706        token_start_line: usize,
707        token_start_column: usize,
708        language: Language,
709    ) -> Result<Option<CodeToken>> {
710        // Check for line comments
711        if let Some(line_comment) = patterns.line_comment {
712            // An empty pattern simply does not match, instead of panicking.
713            if line_comment.starts_with(first_char) {
714                if let Some(token) = self.try_parse_line_comment(
715                    char_indices,
716                    start_offset,
717                    line_comment,
718                    token_start_line,
719                    token_start_column,
720                    language,
721                )? {
722                    return Ok(Some(token));
723                }
724            }
725        }
726
727        // Check for block comments
728        if let Some((start_delim, end_delim)) = patterns.block_comment {
729            // An empty delimiter simply does not match, instead of panicking.
730            if start_delim.starts_with(first_char) {
731                if let Some(token) = self.try_parse_block_comment(
732                    char_indices,
733                    start_offset,
734                    start_delim,
735                    end_delim,
736                    token_start_line,
737                    token_start_column,
738                    language,
739                )? {
740                    return Ok(Some(token));
741                }
742            }
743        }
744
745        Ok(None)
746    }
747
748    /// Parse a line comment
749    fn try_parse_line_comment(
750        &self,
751        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
752        start_offset: usize,
753        comment_start: &str,
754        token_start_line: usize,
755        token_start_column: usize,
756        language: Language,
757    ) -> Result<Option<CodeToken>> {
758        let mut text = String::new();
759        text.push_str(comment_start);
760
761        // Skip the remaining characters of the comment start
762        for _ in 1..comment_start.len() {
763            if let Some((_, ch)) = char_indices.next() {
764                text.push(ch);
765            }
766        }
767
768        // Read until end of line
769        let mut end_offset = start_offset;
770        while let Some((offset, ch)) = char_indices.peek() {
771            if *ch == '\n' {
772                break;
773            }
774            text.push(*ch);
775            end_offset = *offset;
776            char_indices.next();
777        }
778
779        Ok(Some(CodeToken {
780            text,
781            token_type: CodeTokenType::Comment,
782            position: TokenPosition {
783                line: token_start_line,
784                column: token_start_column,
785                start_offset,
786                end_offset,
787            },
788            language,
789        }))
790    }
791
792    /// Parse a block comment
793    fn try_parse_block_comment(
794        &self,
795        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
796        start_offset: usize,
797        start_delim: &str,
798        end_delim: &str,
799        token_start_line: usize,
800        token_start_column: usize,
801        language: Language,
802    ) -> Result<Option<CodeToken>> {
803        let mut text = String::new();
804        text.push_str(start_delim);
805
806        // Skip the remaining characters of the start delimiter
807        for _ in 1..start_delim.len() {
808            if let Some((_, ch)) = char_indices.next() {
809                text.push(ch);
810            }
811        }
812
813        // Read until end delimiter
814        let mut end_offset = start_offset;
815        let end_chars: Vec<char> = end_delim.chars().collect();
816        let mut buffer = Vec::new();
817
818        for (offset, ch) in char_indices.by_ref() {
819            text.push(ch);
820            end_offset = offset;
821            buffer.push(ch);
822
823            // Keep only the last few characters needed to match end delimiter
824            if buffer.len() > end_chars.len() {
825                buffer.remove(0);
826            }
827
828            // Check if we've found the end delimiter
829            if buffer.len() == end_chars.len() && buffer == end_chars {
830                break;
831            }
832        }
833
834        Ok(Some(CodeToken {
835            text,
836            token_type: CodeTokenType::Comment,
837            position: TokenPosition {
838                line: token_start_line,
839                column: token_start_column,
840                start_offset,
841                end_offset,
842            },
843            language,
844        }))
845    }
846
847    /// Parse a string literal
848    fn parse_string_literal(
849        &self,
850        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
851        start_offset: usize,
852        quote_char: char,
853        token_start_line: usize,
854        token_start_column: usize,
855        language: Language,
856    ) -> Result<CodeToken> {
857        let mut text = String::new();
858        text.push(quote_char);
859        let mut end_offset = start_offset;
860        let mut escaped = false;
861
862        for (offset, ch) in char_indices.by_ref() {
863            text.push(ch);
864            end_offset = offset;
865
866            if escaped {
867                escaped = false;
868                continue;
869            }
870
871            if ch == '\\' {
872                escaped = true;
873                continue;
874            }
875
876            if ch == quote_char {
877                break;
878            }
879        }
880
881        Ok(CodeToken {
882            text,
883            token_type: CodeTokenType::String,
884            position: TokenPosition {
885                line: token_start_line,
886                column: token_start_column,
887                start_offset,
888                end_offset,
889            },
890            language,
891        })
892    }
893
894    /// Parse a numeric literal
895    fn parse_numeric_literal(
896        &self,
897        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
898        start_offset: usize,
899        first_char: char,
900        token_start_line: usize,
901        token_start_column: usize,
902        language: Language,
903    ) -> Result<CodeToken> {
904        let mut text = String::new();
905        text.push(first_char);
906        let mut end_offset = start_offset;
907        let mut has_dot = first_char == '.';
908
909        while let Some((offset, ch)) = char_indices.peek() {
910            if ch.is_ascii_digit()
911                || (*ch == '.' && !has_dot)
912                || (*ch == 'e' || *ch == 'E')
913                || (*ch == 'x' || *ch == 'X')
914                || (*ch == '_')
915                || ch.is_ascii_hexdigit()
916            {
917                if *ch == '.' {
918                    has_dot = true;
919                }
920                text.push(*ch);
921                end_offset = *offset;
922                char_indices.next();
923            } else {
924                break;
925            }
926        }
927
928        Ok(CodeToken {
929            text,
930            token_type: CodeTokenType::Number,
931            position: TokenPosition {
932                line: token_start_line,
933                column: token_start_column,
934                start_offset,
935                end_offset,
936            },
937            language,
938        })
939    }
940
941    /// Parse an identifier or keyword
942    fn parse_identifier(
943        &self,
944        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
945        start_offset: usize,
946        first_char: char,
947        token_start_line: usize,
948        token_start_column: usize,
949        language: Language,
950    ) -> Result<CodeToken> {
951        let mut text = String::new();
952        text.push(first_char);
953        let mut end_offset = start_offset;
954
955        while let Some((offset, ch)) = char_indices.peek() {
956            if ch.is_alphanumeric() || *ch == '_' || *ch == '$' {
957                text.push(*ch);
958                end_offset = *offset;
959                char_indices.next();
960            } else {
961                break;
962            }
963        }
964
965        let token_type = if self.keywords.contains(&text) {
966            CodeTokenType::Keyword
967        } else {
968            CodeTokenType::Identifier
969        };
970
971        Ok(CodeToken {
972            text,
973            token_type,
974            position: TokenPosition {
975                line: token_start_line,
976                column: token_start_column,
977                start_offset,
978                end_offset,
979            },
980            language,
981        })
982    }
983
984    /// Parse an operator or punctuation
985    fn parse_operator_or_punctuation(
986        &self,
987        char_indices: &mut std::iter::Peekable<std::str::CharIndices>,
988        start_offset: usize,
989        first_char: char,
990        token_start_line: usize,
991        token_start_column: usize,
992        language: Language,
993    ) -> Result<CodeToken> {
994        let mut text = String::new();
995        text.push(first_char);
996        let mut end_offset = start_offset;
997
998        // Try to form multi-character operators
999        let operators = [
1000            "==", "!=", "<=", ">=", "&&", "||", "++", "--", "+=", "-=", "*=", "/=", "%=", "<<",
1001            ">>", "::", "->", "=>", "**", "//", "...", "..", ":=", "<=>",
1002        ];
1003
1004        for op in &operators {
1005            if op.starts_with(first_char) && op.len() > 1 {
1006                let chars = op.chars().skip(1);
1007                let mut matched = true;
1008                let mut lookahead = Vec::new();
1009
1010                for expected_char in chars {
1011                    if let Some((offset, ch)) = char_indices.peek() {
1012                        if *ch == expected_char {
1013                            lookahead.push((*offset, *ch));
1014                            char_indices.next();
1015                        } else {
1016                            matched = false;
1017                            break;
1018                        }
1019                    } else {
1020                        matched = false;
1021                        break;
1022                    }
1023                }
1024
1025                if matched {
1026                    text = op.to_string();
1027                    if let Some((offset, _)) = lookahead.last() {
1028                        end_offset = *offset;
1029                    }
1030                    break;
1031                } else {
1032                    // Put back the consumed characters
1033                    for (_, _ch) in lookahead.into_iter().rev() {
1034                        // Note: This is a simplified approach. In a real implementation,
1035                        // you'd need a more sophisticated way to put back characters.
1036                    }
1037                }
1038            }
1039        }
1040
1041        let token_type = match first_char {
1042            '(' | ')' | '[' | ']' | '{' | '}' | ';' | ',' | '.' | ':' => CodeTokenType::Punctuation,
1043            _ => CodeTokenType::Operator,
1044        };
1045
1046        Ok(CodeToken {
1047            text,
1048            token_type,
1049            position: TokenPosition {
1050                line: token_start_line,
1051                column: token_start_column,
1052                start_offset,
1053                end_offset,
1054            },
1055            language,
1056        })
1057    }
1058
1059    /// Get vocabulary size
1060    pub fn vocab_size(&self) -> usize {
1061        self.token_to_id.len()
1062    }
1063
1064    /// Get token ID
1065    pub fn token_to_id(&self, token: &str) -> Option<u32> {
1066        self.token_to_id.get(token).copied()
1067    }
1068
1069    /// Get token from ID
1070    pub fn id_to_token(&self, id: u32) -> Option<String> {
1071        self.id_to_token.get(&id).cloned()
1072    }
1073}
1074
1075impl Tokenizer for CodeTokenizer {
1076    fn encode(&self, text: &str) -> Result<TokenizedInput> {
1077        let code_tokens = self.tokenize_code(text)?;
1078        let mut input_ids = Vec::new();
1079
1080        for token in code_tokens {
1081            let token_text = if self.config.normalize_identifiers
1082                && token.token_type == CodeTokenType::Identifier
1083            {
1084                "[IDENTIFIER]".to_string()
1085            } else {
1086                token.text
1087            };
1088
1089            if let Some(id) = self.token_to_id(&token_text) {
1090                input_ids.push(id);
1091            } else if let Some(&unk_id) = self.special_tokens.get("[UNK]") {
1092                input_ids.push(unk_id);
1093            }
1094        }
1095
1096        let attention_mask = vec![1u8; input_ids.len()];
1097
1098        Ok(TokenizedInput {
1099            input_ids,
1100            attention_mask,
1101            token_type_ids: None,
1102            special_tokens_mask: None,
1103            offset_mapping: None,
1104            overflowing_tokens: None,
1105        })
1106    }
1107
1108    fn decode(&self, ids: &[u32]) -> Result<String> {
1109        let tokens: Vec<String> = ids.iter().filter_map(|&id| self.id_to_token(id)).collect();
1110        Ok(tokens.join(" "))
1111    }
1112
1113    fn encode_pair(&self, text_a: &str, text_b: &str) -> Result<TokenizedInput> {
1114        let combined = format!("{}\n{}", text_a, text_b);
1115        self.encode(&combined)
1116    }
1117
1118    fn vocab_size(&self) -> usize {
1119        self.token_to_id.len()
1120    }
1121
1122    fn get_vocab(&self) -> HashMap<String, u32> {
1123        self.token_to_id.clone()
1124    }
1125
1126    fn token_to_id(&self, token: &str) -> Option<u32> {
1127        self.token_to_id.get(token).copied()
1128    }
1129
1130    fn id_to_token(&self, id: u32) -> Option<String> {
1131        self.id_to_token.get(&id).cloned()
1132    }
1133}
1134
1135#[cfg(test)]
1136mod tests {
1137    use super::*;
1138
1139    #[test]
1140    fn test_language_detection() {
1141        assert_eq!(Language::from_extension("rs"), Some(Language::Rust));
1142        assert_eq!(Language::from_extension("py"), Some(Language::Python));
1143        assert_eq!(Language::from_extension("js"), Some(Language::JavaScript));
1144        assert_eq!(Language::from_extension("unknown"), None);
1145    }
1146
1147    #[test]
1148    fn test_rust_tokenization() {
1149        let tokenizer = CodeTokenizer::for_language(Language::Rust);
1150        let code = "fn main() { let x = 42; }";
1151        let tokens = tokenizer.tokenize_code(code).expect("Operation failed in test");
1152
1153        assert!(!tokens.is_empty());
1154
1155        // Check for keywords
1156        let fn_token = tokens.iter().find(|t| t.text == "fn").expect("Operation failed in test");
1157        assert_eq!(fn_token.token_type, CodeTokenType::Keyword);
1158
1159        let let_token = tokens.iter().find(|t| t.text == "let").expect("Operation failed in test");
1160        assert_eq!(let_token.token_type, CodeTokenType::Keyword);
1161    }
1162
1163    #[test]
1164    fn test_string_literal_parsing() {
1165        let tokenizer = CodeTokenizer::for_language(Language::JavaScript);
1166        let code = r#"let name = "Hello \"World\"";"#;
1167        let tokens = tokenizer.tokenize_code(code).expect("Operation failed in test");
1168
1169        let string_token = tokens
1170            .iter()
1171            .find(|t| t.token_type == CodeTokenType::String)
1172            .expect("Operation failed in test");
1173        assert!(string_token.text.starts_with('"'));
1174        assert!(string_token.text.ends_with('"'));
1175    }
1176
1177    #[test]
1178    fn test_comment_parsing() {
1179        let config = CodeTokenizerConfig {
1180            language: Some(Language::Rust),
1181            preserve_comments: true,
1182            ..Default::default()
1183        };
1184        let tokenizer = CodeTokenizer::new(config);
1185        let code = "// This is a comment\nfn main() {}";
1186        let tokens = tokenizer.tokenize_code(code).expect("Operation failed in test");
1187
1188        let comment_token = tokens
1189            .iter()
1190            .find(|t| t.token_type == CodeTokenType::Comment)
1191            .expect("Operation failed in test");
1192        assert!(comment_token.text.starts_with("//"));
1193    }
1194
1195    #[test]
1196    fn test_numeric_literals() {
1197        let tokenizer = CodeTokenizer::for_language(Language::Python);
1198        let code = "x = 42; y = 3.14; z = 0xFF;";
1199        let tokens = tokenizer.tokenize_code(code).expect("Operation failed in test");
1200
1201        let numeric_tokens: Vec<_> =
1202            tokens.iter().filter(|t| t.token_type == CodeTokenType::Number).collect();
1203
1204        assert!(numeric_tokens.len() >= 3);
1205    }
1206
1207    #[test]
1208    fn test_code_tokenizer_encode() {
1209        let tokenizer = CodeTokenizer::for_language(Language::Python);
1210        let code = "def hello(): return 42";
1211        let result = tokenizer.encode(code).expect("Encoding failed");
1212
1213        assert!(!result.input_ids.is_empty());
1214        assert_eq!(result.input_ids.len(), result.attention_mask.len());
1215    }
1216}
trustformers_tokenizers/code_tokenizer.rs

trustformers_tokenizers/
code_tokenizer.rs
