Repository: servo/servo
File: components/webdriver_server/lib.rs (Rust)

Code before:
        let browsing_context_id = wait_for_ipc_response_flatten(receiver)?;
        self.session_mut()?
            .set_browsing_context_id(browsing_context_id);
        // Step 4. Update any implementation-specific state that would result from
        // the user selecting session's current browsing context for interaction,
        // without altering OS-level focus.
        self.focus_browsing_context(browsing_context_id)?;
        Ok(WebDriverResponse::Void)
    }

    /// <https://w3c.github.io/webdriver/#find-element>
    fn handle_find_element(
        &self,
        parameters: &LocatorParameters,
    ) -> WebDriverResult<WebDriverResponse> {
        // Step 1 - 9.
        let res = self.handle_find_elements(parameters)?;
        // Step 10. If result is empty, return error with error code no such element.
        // Otherwise, return the first element of result.
        unwrap_first_element_response(res)
    }

    fn implicit_wait<T>(
        &self,
        callback: impl Fn() -> Result<T, WebDriverError>,
        can_early_return: impl Fn(&T) -> bool,
    ) -> Result<T, WebDriverError>
    {
        let now = Instant::now();
        let (implicit_wait, sleep_interval) = {
            let timeouts = self.session()?.session_timeouts();
            (
                Duration::from_millis(timeouts.implicit_wait.unwrap_or(0)),
                Duration::from_millis(timeouts.sleep_interval),
            )
        };

        loop {
            match callback() {
                Ok(value) if can_early_return(&value) || now.elapsed() > implicit_wait => {
                    return Ok(value);
                },
                Ok(_) => {},
                Err(error) => return Err(error),
            }
            sleep(sleep_interval);
        }
    }

    /// <https://w3c.github.io/webdriver/#find-elements>
    fn handle_find_elements(

Review comment (refactor):
Hrm. Another option here to simplify things is to make it so that callback is `callback: impl Fn() -> Result<Option<T>, WebDriverError>` and ignore `None` results. That should give you more flexibility with how you implement the `callback` because it could also depend on some other kind of state in a more straight-forward way.

Code after:
        let browsing_context_id = wait_for_ipc_response_flatten(receiver)?;
        self.session_mut()?
            .set_browsing_context_id(browsing_context_id);
        // Step 4. Update any implementation-specific state that would result from
        // the user selecting session's current browsing context for interaction,
        // without altering OS-level focus.
        self.focus_browsing_context(browsing_context_id)?;
        Ok(WebDriverResponse::Void)
    }

    /// <https://w3c.github.io/webdriver/#find-element>
    fn handle_find_element(
        &self,
        parameters: &LocatorParameters,
    ) -> WebDriverResult<WebDriverResponse> {
        // Step 1 - 9.
        let res = self.handle_find_elements(parameters)?;
        // Step 10. If result is empty, return error with error code no such element.
        // Otherwise, return the first element of result.
        unwrap_first_element_response(res)
    }

    /// `callback` bound: when it returns Ok, the first value indicates whether
    /// implicit_wait can early return.
    fn implicit_wait<T>(
        &self,
        callback: impl Fn() -> Result<(bool, T), WebDriverError>,
    ) -> Result<T, WebDriverError> {
        let now = Instant::now();
        let (implicit_wait, sleep_interval) = {
            let timeouts = self.session()?.session_timeouts();
            (
                Duration::from_millis(timeouts.implicit_wait.unwrap_or(0)),
                Duration::from_millis(timeouts.sleep_interval),
            )
        };

        loop {
            match callback() {
                Ok((can_early_return, value)) => {
                    if can_early_return || now.elapsed() >= implicit_wait {
                        return Ok(value);
                    }
                },
                Err(error) => return Err(error),
            }
            sleep(sleep_interval);
        }
    }

    /// <https://w3c.github.io/webdriver/#find-elements>