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}