lib.rs

1//! Experimental move codec. PGN lexical/header policy remains Python-owned.
2//! No changed token IDs, no engine evaluation, no inference probability masking.
3use cozy_chess::{Board, Move, Piece, Color};
4use cozy_chess::util::{parse_san_move,parse_uci_move,display_uci_move,display_san_move};
5
6pub const BOS:u8=69;
7pub const EOS:u8=70;
8pub struct Encoded {pub tokens:Vec<u8>,pub uci:Vec<String>,pub fen:String}
9
10fn normalized_fen(board:&Board,halfmove:u32,fullmove:u32)->String {
11    let raw=board.to_string();
12    let mut parts:Vec<String>=raw.split_whitespace().map(str::to_string).collect();
13    if board.en_passant().is_some() {
14        let mut legal_ep=false;
15        board.generate_moves(|set| {
16            for mv in set {if board.piece_on(mv.from)==Some(Piece::Pawn)
17                && mv.from.file()!=mv.to.file() && board.piece_on(mv.to).is_none(){legal_ep=true;}}
18            false
19        });
20        if !legal_ep {parts[3]="-".into();}
21    }
22    parts[4]=halfmove.to_string();parts[5]=fullmove.to_string();parts.join(" ")
23}
24fn initial(fen:Option<&str>)->Result<(Board,u32,u32),String> {
25    if let Some(f)=fen {
26        let p:Vec<_>=f.split_whitespace().collect();
27        if p.len()!=6{return Err("invalid_fen".into())}
28        if p[2].chars().any(|c|!"KQkq-".contains(c)){return Err("unsupported_castling_variant".into())}
29        let half=p[4].parse::<u32>().map_err(|_|"invalid_halfmove")?;
30        let full=p[5].parse::<u32>().map_err(|_|"invalid_fullmove")?;
31        if full==0{return Err("invalid_fullmove".into())}
32        let neutral=format!("{} {} {} {} 0 1",p[0],p[1],p[2],p[3]);
33        let board=Board::from_fen(&neutral,false).map_err(|_|"invalid_fen")?;
34        Ok((board,half,full))
35    } else {Ok((Board::default(),0,1))}
36}
37fn apply(board:&mut Board,mv:Move,half:&mut u32,full:&mut u32)->Result<(String,[u8;3]),String> {
38    if !board.is_legal(mv){return Err("illegal_move".into())}
39    let uci=display_uci_move(board,mv).to_string();let bytes=uci.as_bytes();
40    let from=(bytes[0]-b'a')+8*(bytes[1]-b'1');let to=(bytes[2]-b'a')+8*(bytes[3]-b'1');
41    let promo=match mv.promotion {None=>64,Some(Piece::Knight)=>65,Some(Piece::Bishop)=>66,
42        Some(Piece::Rook)=>67,Some(Piece::Queen)=>68,_=>return Err("invalid_promotion".into())};
43    let pawn=board.piece_on(mv.from)==Some(Piece::Pawn);
44    let capture=board.color_on(mv.to)==Some(!board.side_to_move());
45    *half=if pawn||capture {0}else{half.checked_add(1).ok_or("clock_overflow")?};
46    if board.side_to_move()==Color::Black {*full=full.checked_add(1).ok_or("clock_overflow")?;}
47    board.play(mv);Ok((uci,[from,to,promo]))
48}
49pub fn encode_san(moves:&str,fen:Option<&str>)->Result<Encoded,String> {
50    let (mut board,mut half,mut full)=initial(fen)?;
51    let mut tokens=vec![BOS];let mut uci=Vec::new();
52    for san in moves.split_whitespace() {
53        let normalized=san.trim_end_matches(|c|c=='!'||c=='?').replace('0',"O");
54        let mv=parse_san_move(&board,&normalized).map_err(|_|format!("illegal_san:{san}"))?;
55        let canonical=display_san_move(&board,mv).to_string();
56        if normalized.trim_end_matches(|c|c=='+'||c=='#')!=canonical.trim_end_matches(|c|c=='+'||c=='#')
57            || ((normalized.ends_with('+')||normalized.ends_with('#')) && normalized.chars().last()!=canonical.chars().last()) {
58            return Err("noncanonical_or_false_check_san".into())
59        }
60        let (text,triple)=apply(&mut board,mv,&mut half,&mut full)?;
61        uci.push(text);tokens.extend_from_slice(&triple);
62    }
63    tokens.push(EOS);Ok(Encoded{tokens,uci,fen:normalized_fen(&board,half,full)})
64}
65pub fn decode_tokens(tokens:&[u8],fen:Option<&str>)->Result<Encoded,String> {
66    if tokens.first()!=Some(&BOS)||tokens.last()!=Some(&EOS)||tokens.len()<2||(tokens.len()-2)%3!=0 {
67        return Err("invalid_boundary_or_phase".into())
68    }
69    let (mut board,mut half,mut full)=initial(fen)?;let mut uci=Vec::new();
70    for triple in tokens[1..tokens.len()-1].chunks_exact(3) {
71        let (from,to,promo)=(triple[0],triple[1],triple[2]);
72        if from>=64||to>=64||!(64..=68).contains(&promo){return Err("invalid_token_or_phase".into())}
73        let mut text=format!("{}{}{}{}",(b'a'+from%8)as char,(b'1'+from/8)as char,(b'a'+to%8)as char,(b'1'+to/8)as char);
74        if promo!=64{text.push([b'n',b'b',b'r',b'q'][(promo-65)as usize]as char);}
75        let mv=parse_uci_move(&board,&text).map_err(|_|"invalid_uci")?;
76        let (canonical,triple2)=apply(&mut board,mv,&mut half,&mut full)?;
77        if triple!=triple2{return Err("noncanonical_move".into())}
78        uci.push(canonical);
79    }
80    Ok(Encoded{tokens:tokens.to_vec(),uci,fen:normalized_fen(&board,half,full)})
81}
82
83/// Strict movetext interpreter after the unchanged Python lexical/header gate.
84/// Every variation is validated; only mainline moves enter the model sequence.
85pub fn encode_movetext(text:&str)->Result<(Encoded,String),String> {
86    // Persistent undo chain: moves/branches retain only O(moves+depth) states.
87    // A saved branch copies exactly one undo entry, matching Board.copy(stack=1).
88    // Use an index arena rather than recursive Rc nodes so dropping a very long
89    // legal history cannot overflow the native call stack.
90    struct Node {board:Board,half:u32,full:u32,previous:Option<usize>}
91    let mut nodes=vec![Node{board:Board::default(),half:0,full:1,previous:None}];
92    let mut current=0usize;let mut variations:Vec<usize>=Vec::new();
93    let mut tokens=vec![BOS];let mut uci=Vec::new();let mut ending=None;
94    for atom in text.split_whitespace() {
95        if ending.is_some(){return Err("trailing_movetext".into())}
96        if atom=="@invalid_san_or_null" {return Err("invalid_san_or_null".into())}
97        if atom=="(" {
98            let previous=nodes[current].previous.ok_or("invalid_variation")?;
99            let undo=Node{board:nodes[previous].board.clone(),half:nodes[previous].half,full:nodes[previous].full,previous:None};
100            let saved=Node{board:nodes[current].board.clone(),half:nodes[current].half,full:nodes[current].full,previous:Some(nodes.len())};
101            nodes.push(undo);nodes.push(saved);variations.push(nodes.len()-1);current=previous;continue;
102        }
103        if atom==")" {current=variations.pop().ok_or("unbalanced_variation")?;continue;}
104        if ["1-0","0-1","1/2-1/2","*"].contains(&atom) {
105            if !variations.is_empty(){return Err("result_inside_variation".into())}
106            ending=Some(atom.to_string());continue;
107        }
108        if (atom.starts_with('$')&&atom.len()>1&&atom[1..].bytes().all(|c|c.is_ascii_digit()))
109            || ["!","?","!!","??","!?","?!"].contains(&atom) {continue;}
110        let state=&nodes[current];let (board,half,full)=(&state.board,&state.half,&state.full);
111        if atom.ends_with('.') {
112            let digits=atom.trim_end_matches('.');
113            if !digits.is_empty()&&digits.bytes().all(|c|c.is_ascii_digit()) {
114                if digits.len()>10||digits.parse::<u32>().ok()!=Some(*full)
115                    || (atom.ends_with("...")!=(board.side_to_move()==Color::Black)) {return Err("incorrect_move_number".into())}
116                continue;
117            }
118        }
119        let normalized=atom.trim_end_matches(|c|c=='!'||c=='?').replace('0',"O");
120        let mv=parse_san_move(board,&normalized).map_err(|_|"illegal_or_ambiguous_move")?;
121        let canonical=display_san_move(board,mv).to_string();
122        if normalized.trim_end_matches(|c|c=='+'||c=='#')!=canonical.trim_end_matches(|c|c=='+'||c=='#')
123            || ((normalized.ends_with('+')||normalized.ends_with('#'))&&normalized.chars().last()!=canonical.chars().last()) {
124            return Err("noncanonical_or_false_check_san".into())
125        }
126        let (mut next,mut h,mut f)=(board.clone(),*half,*full);
127        let (word,triple)=apply(&mut next,mv,&mut h,&mut f)?;
128        if variations.is_empty(){uci.push(word);tokens.extend_from_slice(&triple);}
129        nodes.push(Node{board:next,half:h,full:f,previous:Some(current)});current=nodes.len()-1;
130    }
131    if !variations.is_empty()||ending.is_none(){return Err("incomplete_game".into())}
132    tokens.push(EOS);let state=&nodes[current];let (board,half,full)=(&state.board,&state.half,&state.full);
133    let encoded=Encoded{tokens,uci,fen:normalized_fen(board,*half,*full)};
134    // Independent token-path replay is mandatory, including final board equality.
135    let decoded=decode_tokens(&encoded.tokens,None)?;
136    if encoded.uci!=decoded.uci||encoded.fen!=decoded.fen{return Err("round_trip_mismatch".into())}
137    Ok((encoded,ending.unwrap()))
138}
139
140