eval_stack/
judge.rs

1use std::{
2    fs,
3    future::Future,
4    io::{BufRead, BufReader, Read},
5    os::unix::process::ExitStatusExt,
6    path::PathBuf,
7    task::Poll,
8    time::Duration,
9};
10
11use anyhow::Result;
12
13use crate::utils::get_memory_usage;
14
15#[derive(Debug, Clone)]
16#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
17#[serde(rename_all = "snake_case", tag = "type")]
18pub enum JudgeStatus {
19    Accepted,
20    WrongAnswer,
21    TimeLimitExceeded,
22    MemoryLimitExceeded,
23    RuntimeError {
24        code: i32,
25        stderr: String,
26    },
27    CompileError {
28        message: String,
29    },
30    SystemError {
31        code: i32,
32        stderr: String,
33        signal: i32,
34    },
35    SegmentFault {
36        code: i32,
37        stderr: String,
38    },
39}
40
41#[derive(Debug, Clone)]
42#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
43#[serde(rename_all = "camelCase")]
44pub struct JudgeResult {
45    pub status: JudgeStatus,
46    pub time_used: Duration,
47    pub memory_used: u64,
48}
49
50impl Default for JudgeResult {
51    fn default() -> Self {
52        Self {
53            status: JudgeStatus::Accepted,
54            time_used: Duration::from_secs(0),
55            memory_used: 0,
56        }
57    }
58}
59
60impl JudgeResult {
61    pub fn is_accepted(&self) -> bool {
62        matches!(self.status, JudgeStatus::Accepted)
63    }
64
65    pub fn is_wrong_answer(&self) -> bool {
66        !self.is_accepted()
67    }
68}
69
70pub struct Judge {
71    pub child: std::process::Child,
72    pub id: u32,
73    pub time_limit: Duration,
74    pub memory_limit: u64,
75    pub instant: tokio::time::Instant,
76    pub memory_used: u64,
77    pub time_used: Duration,
78    pub stdout_file: PathBuf,
79    pub expected_output_file: PathBuf,
80}
81
82impl Future for Judge {
83    type Output = Result<JudgeResult>;
84
85    fn poll(
86        mut self: std::pin::Pin<&mut Self>,
87        cx: &mut std::task::Context<'_>,
88    ) -> std::task::Poll<Self::Output> {
89        match self.child.try_wait()? {
90            Some(status) => {
91                self.time_used = self.instant.elapsed();
92                drop(self.child.stdin.take());
93                drop(self.child.stdout.take());
94                if status.success() {
95                    let stdout = BufReader::new(fs::File::open(&self.stdout_file)?);
96                    let expected_out = BufReader::new(fs::File::open(&self.expected_output_file)?);
97
98                    let mut stdout_lines = stdout.lines();
99                    let mut expected_out_lines = expected_out.lines();
100
101                    let matched = loop {
102                        match (stdout_lines.next(), expected_out_lines.next()) {
103                            (None, None) => break true,
104                            (Some(output), None) => {
105                                if output?
106                                    .trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
107                                    != ""
108                                {
109                                    break false;
110                                }
111                            }
112                            (None, Some(expected_output)) => {
113                                if expected_output?
114                                    .trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
115                                    != ""
116                                {
117                                    break false;
118                                }
119                            }
120                            (Some(output), Some(expected_output)) => {
121                                if output?
122                                    .trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
123                                    != expected_output?
124                                        .trim_end_matches(|c: char| c.is_whitespace() || c == '\n')
125                                {
126                                    break false;
127                                }
128                            }
129                        }
130                    };
131
132                    if matched {
133                        Poll::Ready(Ok(JudgeResult {
134                            status: JudgeStatus::Accepted,
135                            time_used: self.time_used,
136                            memory_used: self.memory_used,
137                        }))
138                    } else {
139                        Poll::Ready(Ok(JudgeResult {
140                            status: JudgeStatus::WrongAnswer,
141                            time_used: self.time_used,
142                            memory_used: self.memory_used,
143                        }))
144                    }
145                } else {
146                    let mut stderr = String::new();
147                    let _ = self
148                        .child
149                        .stderr
150                        .take()
151                        .unwrap()
152                        .read_to_string(&mut stderr);
153                    let code = status.code().unwrap_or(-1);
154                    match status.signal() {
155                        Some(libc::SIGSEGV) | Some(libc::SIGBUS) | Some(libc::SIGILL) => {
156                            Poll::Ready(Ok(JudgeResult {
157                                status: JudgeStatus::SegmentFault { code, stderr },
158                                time_used: self.time_used,
159                                memory_used: self.memory_used,
160                            }))
161                        }
162                        Some(signal) => Poll::Ready(Ok(JudgeResult {
163                            status: JudgeStatus::SystemError {
164                                code,
165                                signal,
166                                stderr,
167                            },
168                            time_used: self.time_used,
169                            memory_used: self.memory_used,
170                        })),
171                        None => Poll::Ready(Ok(JudgeResult {
172                            status: JudgeStatus::RuntimeError { code, stderr },
173                            time_used: self.time_used,
174                            memory_used: self.memory_used,
175                        })),
176                    }
177                }
178            }
179            None => {
180                if let Some(memory_used) = get_memory_usage(self.id) {
181                    self.memory_used = memory_used.max(self.memory_used);
182                };
183                if self.memory_used > self.memory_limit {
184                    self.child.kill()?;
185                    self.time_used = self.instant.elapsed();
186                    return Poll::Ready(Ok(JudgeResult {
187                        status: JudgeStatus::MemoryLimitExceeded,
188                        time_used: self.time_used,
189                        memory_used: self.memory_used,
190                    }));
191                }
192                if self.instant.elapsed() > self.time_limit {
193                    self.child.kill()?;
194                    self.time_used = self.instant.elapsed();
195                    return Poll::Ready(Ok(JudgeResult {
196                        status: JudgeStatus::TimeLimitExceeded,
197                        time_used: self.time_used,
198                        memory_used: self.memory_used,
199                    }));
200                }
201                cx.waker().wake_by_ref();
202                Poll::Pending
203            }
204        }
205    }
206}