opentitanlib/io/console/
ext.rs

1// Copyright lowRISC contributors (OpenTitan project).
2// Licensed under the Apache License, Version 2.0, see LICENSE for details.
3// SPDX-License-Identifier: Apache-2.0
4
5use std::task::{Poll, ready};
6use std::time::Duration;
7
8use anyhow::{Context, Result};
9
10use super::ConsoleDevice;
11use crate::io::console::{ConsoleError, CoverageMiddleware, Logged};
12
13/// Extension trait to [`Uart`] where useful methods are provided.
14pub trait ConsoleExt {
15    /// Reads UART receive data into `buf`, returning the number of bytes read.
16    /// This function is blocking.
17    fn read(&self, buf: &mut [u8]) -> Result<usize>;
18
19    /// Reads UART receive data into `buf`, returning the number of bytes read.
20    /// The `timeout` may be used to specify a duration to wait for data.
21    /// If timeout expires without any data arriving `Ok(0)` will be returned, never `Err(_)`.
22    fn read_timeout(&self, buf: &mut [u8], timeout: Duration) -> Result<usize>;
23
24    /// Return a wrapper that logs all outputs while reading.
25    fn logged(self) -> Logged<Self>
26    where
27        Self: Sized;
28
29    /// Return a wrapper that extracts coverage data.
30    fn coverage(self) -> CoverageMiddleware<Self>
31    where
32        Self: Sized;
33
34    /// Wait for a line that matches the specified pattern to appear.
35    ///
36    /// The pattern matched is returned. If timeout occurs, `None` is returned.
37    fn try_wait_for_line<P: MatchPattern>(
38        &self,
39        pattern: P,
40        timeout: Duration,
41    ) -> Result<Option<P::MatchResult>>;
42
43    /// Wait for a line that matches the specified pattern to appear.
44    ///
45    /// Types that can be used include:
46    /// * `str`` / `[u8]`: literal matching is performed, no return value
47    /// * `Regex`: regex captures are returned.
48    /// * `(T, E)`, where `T` and `E` are one of the above: match two patterns at once.
49    ///   If the first matches, `Ok` is returned, otherwise `Err` is. Note that when this
50    ///   is used, you would have `anyhow::Result<Result<TMatch, EMatch>>` from this function,
51    ///   where the `Err(_)` is I/O error or timeout, and `Ok(Err(_))` is the match for `E`.
52    ///
53    /// If you want to construct a static `&Regex` you can use [`regex!`] macro.
54    ///
55    /// [`regex!`]: crate::regex
56    ///
57    /// The pattern matched is returned.
58    fn wait_for_line<P: MatchPattern>(
59        &self,
60        pattern: P,
61        timeout: Duration,
62    ) -> Result<P::MatchResult>;
63
64    /// Wait on the console until the coverage profile end or skip anchor is received.
65    fn wait_for_coverage(&self, timeout: Duration) -> Result<()>;
66}
67
68impl<T: ConsoleDevice + ?Sized> ConsoleExt for T {
69    fn read(&self, buf: &mut [u8]) -> Result<usize> {
70        crate::util::runtime::block_on(std::future::poll_fn(|cx| self.poll_read(cx, buf)))
71    }
72
73    fn read_timeout(&self, buf: &mut [u8], timeout: Duration) -> Result<usize> {
74        crate::util::runtime::block_on(async {
75            tokio::time::timeout(timeout, std::future::poll_fn(|cx| self.poll_read(cx, buf))).await
76        })
77        .unwrap_or(Ok(0))
78    }
79
80    fn logged(self) -> Logged<Self>
81    where
82        Self: Sized,
83    {
84        Logged::new(self)
85    }
86
87    fn coverage(self) -> CoverageMiddleware<Self>
88    where
89        Self: Sized,
90    {
91        CoverageMiddleware::new(self)
92    }
93
94    fn try_wait_for_line<P: MatchPattern>(
95        &self,
96        pattern: P,
97        timeout: Duration,
98    ) -> Result<Option<P::MatchResult>> {
99        crate::util::runtime::block_on(async {
100            match tokio::time::timeout(timeout, async {
101                loop {
102                    let line = read_line(self).await?;
103                    if let Some(m) = pattern.perform_match(&line) {
104                        return Ok(m);
105                    }
106                }
107            })
108            .await
109            {
110                Ok(Ok(v)) => Ok(Some(v)),
111                Ok(Err(e)) => Err(e),
112                Err(_) => Ok(None),
113            }
114        })
115    }
116
117    fn wait_for_line<P: MatchPattern>(
118        &self,
119        pattern: P,
120        timeout: Duration,
121    ) -> Result<P::MatchResult> {
122        self.try_wait_for_line(pattern, timeout)?
123            .with_context(|| ConsoleError::TimedOut)
124    }
125
126    fn wait_for_coverage(&self, timeout: Duration) -> Result<()> {
127        let coverage = self.as_coverage_console().ok_or_else(|| {
128            ConsoleError::GenericError("Coverage extraction not supported by this device".into())
129        })?;
130        let initial = coverage.coverage_blocks_processed();
131        crate::util::runtime::block_on(async {
132            tokio::time::timeout(timeout, async {
133                std::future::poll_fn(|cx| {
134                    if coverage.coverage_blocks_processed() > initial {
135                        return Poll::Ready(Ok(()));
136                    }
137                    let mut buf = [0u8; 1024];
138                    match ready!(self.poll_read(cx, &mut buf)) {
139                        Ok(0) => Poll::Ready(Err(ConsoleError::GenericError(
140                            "EOF reached while waiting for coverage".into(),
141                        )
142                        .into())),
143                        Ok(_) => {
144                            cx.waker().wake_by_ref();
145                            Poll::Pending
146                        }
147                        Err(e) => Poll::Ready(Err(e)),
148                    }
149                })
150                .await
151            })
152            .await
153            .context("Timed out waiting for coverage")?
154        })
155    }
156}
157
158async fn read_line<T: ConsoleDevice + ?Sized>(console: &T) -> Result<Vec<u8>> {
159    let mut buf = Vec::new();
160
161    loop {
162        // Read one byte at a time to avoid the risk of consuming past a line.
163        let mut ch = 0;
164        let len =
165            std::future::poll_fn(|cx| console.poll_read(cx, std::slice::from_mut(&mut ch))).await?;
166        if len == 0 {
167            break;
168        }
169
170        buf.push(ch);
171        if ch == b'\n' {
172            break;
173        }
174    }
175
176    Ok(buf)
177}
178
179/// Indicating types that can be used for `wait_for_line` matching.
180pub trait MatchPattern {
181    type MatchResult;
182
183    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult>;
184}
185
186impl<T: MatchPattern + ?Sized> MatchPattern for &T {
187    type MatchResult = T::MatchResult;
188
189    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
190        T::perform_match(self, haystack)
191    }
192}
193
194impl MatchPattern for [u8] {
195    type MatchResult = ();
196
197    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
198        memchr::memmem::find(haystack, self).map(|_| ())
199    }
200}
201
202impl MatchPattern for regex::bytes::Regex {
203    type MatchResult = Vec<Vec<u8>>;
204
205    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
206        Some(
207            self.captures(haystack)?
208                .iter()
209                .map(|x| x.map(|m| m.as_bytes().to_owned()).unwrap_or_default())
210                .collect(),
211        )
212    }
213}
214
215impl MatchPattern for str {
216    type MatchResult = ();
217
218    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
219        self.as_bytes().perform_match(haystack)
220    }
221}
222
223impl MatchPattern for regex::Regex {
224    type MatchResult = Vec<String>;
225
226    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
227        let haystack = String::from_utf8_lossy(haystack);
228        Some(
229            self.captures(&haystack)?
230                .iter()
231                .map(|x| x.map(|m| m.as_str().to_owned()).unwrap_or_default())
232                .collect(),
233        )
234    }
235}
236
237/// Match two patterns at once.
238pub struct PassFail<T, E>(pub T, pub E);
239
240pub enum PassFailResult<T, E> {
241    Pass(T),
242    Fail(E),
243}
244
245impl<T: MatchPattern, E: MatchPattern> MatchPattern for PassFail<T, E> {
246    type MatchResult = PassFailResult<T::MatchResult, E::MatchResult>;
247
248    fn perform_match(&self, haystack: &[u8]) -> Option<Self::MatchResult> {
249        if let Some(m) = self.1.perform_match(haystack) {
250            return Some(PassFailResult::Fail(m));
251        }
252
253        Some(PassFailResult::Pass(self.0.perform_match(haystack)?))
254    }
255}
opentitanlib/io/console/ext.rs

opentitanlib/io/console/
ext.rs
