main.rs

1use std::{collections::HashMap,io::{self,Read},time::{Duration,Instant}};
2use std::sync::{Arc,Mutex,atomic::{AtomicU64,AtomicU8,AtomicUsize,Ordering}};
3use serde::{Deserialize,Serialize};
4use shakmaty::{Chess,Position,Color,CastlingMode,EnPassantMode,Move,fen::{Fen,Epd},uci::UciMove};
5#[derive(Deserialize)] #[serde(deny_unknown_fields)]
6struct Request { root_fen:String,moves:Vec<String>,max_plies:u8,max_nodes:u64,millis:u64,workers:usize }
7#[derive(Clone,Serialize)]
8struct Branch { r#move:String,child:Proof }
9#[derive(Clone,Serialize)]
10struct Proof { kind:&'static str,fen:String,mate_plies:u8,branches:Vec<Branch> }
11#[derive(Default)]
12struct Stats { nodes:AtomicU64,deepest:AtomicU8 }
13struct Search { attacker:Color,start:Instant,budget:Duration,stats:Arc<Stats>,max_nodes:u64,counts:HashMap<String,u16>,cutoff:Option<(Arc<AtomicUsize>,usize)> }
14fn key(p:&Chess)->String { Epd::from_position(p.clone(),EnPassantMode::Legal).to_string() }
15fn fen(p:&Chess)->String { Fen::from_position(p.clone(),EnPassantMode::Legal).to_string() }
16fn uci(m:&Move)->String { UciMove::from_move(m,CastlingMode::Standard).to_string() }
17impl Search {
18 fn guard(&self)->Result<(),&'static str> {
19  if self.start.elapsed()>=self.budget {return Err("deadline")}
20  if let Some((best,index))=&self.cutoff {if *index>best.load(Ordering::Relaxed){return Err("superseded")}}
21  Ok(())
22 }
23 fn node(&self,ply:u8)->Result<(),&'static str> {
24  self.guard()?;
25  self.stats.nodes.fetch_update(Ordering::Relaxed,Ordering::Relaxed,|n|if n<self.max_nodes{Some(n+1)}else{None}).map_err(|_|"node_budget")?;
26  self.stats.deepest.fetch_max(ply,Ordering::Relaxed);Ok(())
27 }
28 fn root(&mut self,p:&Chess,depth:u8,workers:usize)->Result<Option<Proof>,&'static str> {
29  if workers==1 || depth==1{return self.visit(p,depth,0)}
30  self.node(0)?;let moves=p.legal_moves();
31  if moves.is_empty() || self.draw(p,&moves)? {return Ok(None)}
32  let mut ordered:Vec<_>=moves.iter().map(|m|{let mut q=p.clone();q.play_unchecked(m);(!q.is_check(),uci(m),q)}).collect();
33  ordered.sort_by(|a,b|(&a.0,&a.1).cmp(&(&b.0,&b.1)));
34  let next=AtomicUsize::new(0);let best=Arc::new(AtomicUsize::new(ordered.len()));
35  let results=Mutex::new(vec![None;ordered.len()]);
36  std::thread::scope(|scope| {
37   for _ in 0..workers.min(ordered.len()) {
38    let tasks=&ordered;let next=&next;let results=&results;let best=best.clone();
39    let stats=self.stats.clone();let base=self.counts.clone();let attacker=self.attacker;let start=self.start;let budget=self.budget;let max_nodes=self.max_nodes;
40    scope.spawn(move || loop {
41     let index=next.fetch_add(1,Ordering::Relaxed);
42     if index>=tasks.len() || index>best.load(Ordering::Relaxed){break}
43     let q=&tasks[index].2;let mut local=Search{attacker,start,budget,stats:stats.clone(),max_nodes,counts:base.clone(),cutoff:Some((best.clone(),index))};
44     *local.counts.entry(key(q)).or_default()+=1;
45     let result=local.visit(q,depth-1,1);
46     if matches!(&result,Ok(Some(_))){best.fetch_min(index,Ordering::Relaxed);}
47     results.lock().unwrap()[index]=Some(result);
48    });
49   }
50  });
51  self.guard()?;
52  for (index,result) in results.into_inner().unwrap().into_iter().enumerate() {
53   match result {
54    Some(Ok(Some(proof)))=>return Ok(Some(Proof{kind:"attacker_choice",fen:fen(p),mate_plies:1+proof.mate_plies,branches:vec![Branch{r#move:ordered[index].1.clone(),child:proof}]})),
55    Some(Ok(None))=>(),Some(Err(e))=>return Err(e),None=>return Err("deadline")
56   }
57  }
58  Ok(None)
59 }
60 fn draw(&self,p:&Chess,moves:&[Move])->Result<bool,&'static str> {
61  self.guard()?;
62  if p.is_insufficient_material() || p.halfmoves()>=100 || self.counts.get(&key(p)).copied().unwrap_or(0)>=3 {return Ok(true)}
63  // Same immediate symmetric claim policy as python-chess claim_draw=True,
64  // including declaring a legal move to claim fifty moves or threefold.
65  for m in moves {
66   self.guard()?;let mut q=p.clone();q.play_unchecked(m);
67   if self.counts.get(&key(&q)).copied().unwrap_or(0)>=2 {return Ok(true)}
68   if p.halfmoves()>=99 && q.halfmoves()>=100 && !q.legal_moves().is_empty() {return Ok(true)}
69  }
70  Ok(false)
71 }
72 fn visit(&mut self,p:&Chess,remaining:u8,ply:u8)->Result<Option<Proof>,&'static str> {
73  self.node(ply)?;
74  let moves=p.legal_moves();
75  if moves.is_empty() {return Ok(if p.is_check() && p.turn()!=self.attacker {Some(Proof{kind:"checkmate",fen:fen(p),mate_plies:0,branches:vec![]})}else{None})}
76  if self.draw(p,&moves)? || remaining==0{return Ok(None)}
77  // Checks first is ordering only: quiet moves and every legal defense remain.
78  let mut ordered:Vec<_>=moves.iter().map(|m| {let mut q=p.clone();q.play_unchecked(m);(!q.is_check(),uci(m),q)}).collect();
79  ordered.sort_by(|a,b|(&a.0,&a.1).cmp(&(&b.0,&b.1)));
80  let attacking=p.turn()==self.attacker;let mut branches=vec![];
81  for (_,mv,q) in ordered {
82   self.guard()?;let k=key(&q);*self.counts.entry(k.clone()).or_default()+=1;
83   let child=self.visit(&q,remaining-1,ply+1);
84   let n=self.counts.get_mut(&k).unwrap();*n-=1;if *n==0{self.counts.remove(&k);}
85   self.guard()?;
86   match child? {
87    Some(proof)=>{let branch=Branch{r#move:mv,child:proof};if attacking{return Ok(Some(Proof{kind:"attacker_choice",fen:fen(p),mate_plies:1+branch.child.mate_plies,branches:vec![branch]}))}branches.push(branch);}
88    None=>if !attacking{return Ok(None)},
89   }
90  }
91  Ok(if !attacking {Some(Proof{kind:"all_defenses",fen:fen(p),mate_plies:1+branches.iter().map(|b|b.child.mate_plies).max().unwrap(),branches})}else{None})
92 }
93}
94fn run(r:Request)->Result<serde_json::Value,String> {
95 if !(1..=10).contains(&r.max_plies)||!(1..=8).contains(&r.workers)||r.max_nodes==0||r.max_nodes>1_000_000||r.millis==0||r.millis>10_000||r.moves.len()>1024{return Err("invalid bounded request".into())}
96 let start=Instant::now();let f:Fen=r.root_fen.parse().map_err(|_|"invalid fen")?;
97 let mut p:Chess=f.into_position(CastlingMode::Standard).map_err(|_|"invalid standard position")?;
98 let mut counts=HashMap::new();counts.insert(key(&p),1);
99 for text in r.moves {let m:UciMove=text.parse().map_err(|_|"invalid uci")?;let m=m.to_move(&p).map_err(|_|"illegal history")?;p.play_unchecked(&m);*counts.entry(key(&p)).or_default()+=1;}
100 let mut s=Search{attacker:p.turn(),start,budget:Duration::from_millis(r.millis),stats:Arc::new(Stats::default()),max_nodes:r.max_nodes,counts,cutoff:None};
101 let mut completed=0;let mut proof=None;let mut status="no_forced_mate";let mut reason="exhausted_depth_bound";
102 let mut iterations=vec![];
103 for depth in (1..=r.max_plies).step_by(2) {iterations.push(depth);match s.root(&p,depth,r.workers) {
104  Ok(value)=>{if s.guard().is_err(){status="unknown";reason="deadline";break}completed=depth;if value.is_some(){proof=value;status="proven";reason="forced_checkmate";break}},
105  Err(e)=>{status="unknown";reason=e;break}
106 }}
107 if status=="no_forced_mate" {completed=r.max_plies;} // A root-side mate is possible only at odd plies.
108 let mv=proof.as_ref().and_then(|p|p.branches.first()).map(|b|b.r#move.clone());
109 Ok(serde_json::json!({"status":status,"reason":reason,"workers":r.workers,"iteration_depths":iterations,"max_plies":r.max_plies,"completed_depth_plies":completed,"deepest_visited_ply":s.stats.deepest.load(Ordering::Relaxed),"mate_plies":proof.as_ref().map(|p|p.mate_plies),"move":mv,"nodes":s.stats.nodes.load(Ordering::Relaxed),"node_budget":r.max_nodes,"seconds":s.start.elapsed().as_secs_f64(),"time_budget_seconds":r.millis as f64/1000.,"proof":proof,"backend":"rust-shakmaty-0.27.3","ordering":"checks-first then lexical UCI; all legal moves","transpositions":"disabled; exact repetition counts retained"}))
110}
111fn main() {let mut input=String::new();let result=io::stdin().take(131073).read_to_string(&mut input).map_err(|e|e.to_string()).and_then(|_|if input.len()>131072{Err("request too large".into())}else{serde_json::from_str::<Request>(&input).map_err(|e|e.to_string())}).and_then(run);match result {Ok(v)=>println!("{}",v),Err(e)=>{println!("{}",serde_json::json!({"status":"unknown","reason":"invalid_request","error":e}));std::process::exit(2)}}}
112