AnyBot/src/marks.rs

// What a bot sees of a screen: the picture with every control it can use boxed and numbered,
// and the same numbers in a list ("[12] 連結「Profile」"). It acts by number; the click lands on
// that control's centre. Guessing coordinates from a picture is where clicks went wrong (a
// model took an empty strip for a menu item and clicked it ten times); a number can't miss.
//
// The controls come from the system's accessibility tree: macOS AX and Linux AT-SPI through
// agent-desktop on a person's computer (desktop.rs), AT-SPI in the bots' own computer
// (box/a11y). Web pages are in those trees too (Chromium with renderer accessibility on).

use std::collections::HashMap;

/// One control on the screen, in screen units (the mouse's: points on a Mac, pixels elsewhere).
#[derive(Clone, Debug, PartialEq)]
pub struct Element {
    /// 按鈕, 連結, 輸入框… (already in words).
    pub kind: String,
    pub name: String,
    pub value: String,
    pub x: f64,
    pub y: f64,
    pub w: f64,
    pub h: f64,
    /// Something to click or type into (gets a number), not just text.
    pub active: bool,
    pub flags: Vec<&'static str>,
    pub window: String,
}

impl Element {
    pub fn centre(&self) -> (f64, f64) {
        (self.x + self.w / 2.0, self.y + self.h / 2.0)
    }
}

/// Most controls numbered on one picture (more is clutter the model can't use).
pub const MAX_MARKS: usize = 150;

/// The numbered controls of the latest look at a screen: number → (centre, words).
#[derive(Clone, Debug, Default)]
pub struct Spots {
    pub by_number: HashMap<usize, ((f64, f64), String)>,
}

impl Spots {
    /// "12", "e12", "[12]" → the control's centre and what it is.
    pub fn get(&self, r: &str) -> Result<((f64, f64), String), String> {
        let n = number_of(r).ok_or_else(|| format!("「{}」不是編號：請用截圖上的編號（例如 12）", r.trim()))?;
        self.by_number.get(&n).cloned().ok_or_else(|| format!("最新的畫面裡沒有編號 {n}：先 computer_use_screenshot 再看一次。"))
    }
}

pub fn number_of(r: &str) -> Option<usize> {
    let t = r.trim().trim_matches(['[', ']']).trim_start_matches(['e', 'E', '#']);
    t.parse().ok()
}

/// Order (reading order, by window) and number the controls; the list's lines and the spots.
/// `all_text`: also list plain text (when there is no picture to read it from).
pub fn number(elements: &[Element], all_text: bool) -> (Vec<(usize, Element)>, Vec<String>, Spots) {
    let mut seen = std::collections::HashSet::new();
    let mut els: Vec<&Element> = elements.iter().filter(|e| e.w >= 1.0 && e.h >= 1.0 && seen.insert((e.kind.clone(), e.name.clone(), e.x as i64, e.y as i64, e.w as i64, e.h as i64))).collect();
    let mut windows: Vec<String> = vec![];
    for e in &els {
        if !windows.contains(&e.window) {
            windows.push(e.window.clone());
        }
    }
    els.sort_by(|a, b| {
        let wa = windows.iter().position(|w| *w == a.window).unwrap_or(0);
        let wb = windows.iter().position(|w| *w == b.window).unwrap_or(0);
        (wa, (a.y / 12.0) as i64, a.x as i64).cmp(&(wb, (b.y / 12.0) as i64, b.x as i64))
    });
    let mut numbered = vec![];
    let mut lines = vec![];
    let mut spots = Spots::default();
    let mut n = 0;
    let mut window = None;
    for e in els {
        if !e.active && !all_text {
            continue;
        }
        if windows.len() > 1 && window.as_ref() != Some(&e.window) && !e.window.is_empty() {
            window = Some(e.window.clone());
            lines.push(format!("## {}", e.window));
        }
        let what = describe(e);
        if e.active && n < MAX_MARKS {
            n += 1;
            spots.by_number.insert(n, (e.centre(), what.clone()));
            numbered.push((n, e.clone()));
            lines.push(format!("[{n}] {what}"));
        } else if !e.active {
            lines.push(format!("    {what}"));
        }
    }
    (numbered, lines, spots)
}

fn describe(e: &Element) -> String {
    let short = |s: &str, n: usize| -> String {
        let c: Vec<char> = s.chars().collect();
        if c.len() > n { format!("{}…", c[..n].iter().collect::<String>()) } else { s.to_string() }
    };
    let mut t = e.kind.clone();
    if !e.name.is_empty() {
        t.push_str(&format!("「{}」", short(&e.name, 60)));
    }
    if !e.value.is_empty() && e.value != e.name {
        t.push_str(&format!("：{}", short(&e.value, 80)));
    }
    if !e.flags.is_empty() {
        t.push_str(&format!("（{}）", e.flags.join("、")));
    }
    t
}

/// A fingerprint of what the controls are and where (to say "nothing changed").
pub fn fingerprint(elements: &[Element]) -> u64 {
    use std::hash::{Hash, Hasher};
    let mut h = std::collections::hash_map::DefaultHasher::new();
    for e in elements {
        (e.kind.as_str(), e.name.as_str(), e.value.as_str(), e.x as i64, e.y as i64, e.w as i64, e.h as i64, e.flags.len()).hash(&mut h);
    }
    h.finish()
}

/// A coarse fingerprint of a picture (16×10 grey cells), for "nothing changed" when there is no
/// tree to compare.
pub fn picture_fingerprint(img: &image::RgbImage) -> u64 {
    use std::hash::{Hash, Hasher};
    let small = image::imageops::resize(img, 16, 10, image::imageops::FilterType::Triangle);
    let mut h = std::collections::hash_map::DefaultHasher::new();
    for p in small.pixels() {
        // Quantised, so a blinking caret or a clock's second doesn't count as a change.
        ((p[0] as u16 + p[1] as u16 + p[2] as u16) / 24).hash(&mut h);
    }
    h.finish()
}

const COLOURS: [[u8; 3]; 6] = [[220, 38, 38], [37, 99, 235], [22, 163, 74], [192, 38, 211], [234, 88, 12], [8, 145, 178]];

/// Box and number each control on the picture. `to_pic` maps screen units to picture pixels.
pub fn draw(img: &mut image::RgbImage, numbered: &[(usize, Element)], to_pic: (f64, f64)) {
    let (iw, ih) = (img.width() as i64, img.height() as i64);
    for (n, e) in numbered {
        let c = COLOURS[n % COLOURS.len()];
        let (x0, y0) = ((e.x / to_pic.0).round() as i64, (e.y / to_pic.1).round() as i64);
        let (x1, y1) = (((e.x + e.w) / to_pic.0).round() as i64, ((e.y + e.h) / to_pic.1).round() as i64);
        rect_outline(img, x0, y0, x1, y1, c);
        // The number: white digits on the control's colour, at its top-left corner (inside the
        // picture).
        let label = n.to_string();
        let (lw, lh) = (label.len() as i64 * (DIGIT_W + 1) * SCALE + 3, DIGIT_H * SCALE + 4);
        let lx = x0.clamp(0, (iw - lw).max(0));
        let ly = (y0 - lh).clamp(0, (ih - lh).max(0));
        let ly = if ly + lh > y0 && y0 >= lh { y0 - lh } else { ly };
        fill(img, lx, ly, lx + lw, ly + lh, c);
        let mut cx = lx + 2;
        for ch in label.bytes() {
            digit(img, cx, ly + 2, ch - b'0', [255, 255, 255]);
            cx += (DIGIT_W + 1) * SCALE;
        }
    }
}

fn put(img: &mut image::RgbImage, x: i64, y: i64, c: [u8; 3]) {
    if x >= 0 && y >= 0 && (x as u32) < img.width() && (y as u32) < img.height() {
        img.put_pixel(x as u32, y as u32, image::Rgb(c));
    }
}

fn fill(img: &mut image::RgbImage, x0: i64, y0: i64, x1: i64, y1: i64, c: [u8; 3]) {
    for y in y0..y1 {
        for x in x0..x1 {
            put(img, x, y, c);
        }
    }
}

fn rect_outline(img: &mut image::RgbImage, x0: i64, y0: i64, x1: i64, y1: i64, c: [u8; 3]) {
    for x in x0..=x1 {
        put(img, x, y0, c);
        put(img, x, y1, c);
    }
    for y in y0..=y1 {
        put(img, x0, y, c);
        put(img, x1, y, c);
    }
}

/// A 3×5 digit font, each pixel drawn SCALE×SCALE.
const DIGIT_W: i64 = 3;
const DIGIT_H: i64 = 5;
const SCALE: i64 = 2;
const DIGITS: [[u8; 5]; 10] = [
    [0b111, 0b101, 0b101, 0b101, 0b111],
    [0b010, 0b110, 0b010, 0b010, 0b111],
    [0b111, 0b001, 0b111, 0b100, 0b111],
    [0b111, 0b001, 0b111, 0b001, 0b111],
    [0b101, 0b101, 0b111, 0b001, 0b001],
    [0b111, 0b100, 0b111, 0b001, 0b111],
    [0b111, 0b100, 0b111, 0b101, 0b111],
    [0b111, 0b001, 0b010, 0b010, 0b010],
    [0b111, 0b101, 0b111, 0b101, 0b111],
    [0b111, 0b101, 0b111, 0b001, 0b111],
];

fn digit(img: &mut image::RgbImage, x: i64, y: i64, d: u8, c: [u8; 3]) {
    let rows = DIGITS[(d as usize).min(9)];
    for (r, bits) in rows.iter().enumerate() {
        for col in 0..DIGIT_W {
            if bits >> (DIGIT_W - 1 - col) & 1 == 1 {
                fill(img, x + col * SCALE, y + r as i64 * SCALE, x + (col + 1) * SCALE, y + (r as i64 + 1) * SCALE, c);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn el(kind: &str, name: &str, x: f64, y: f64, active: bool) -> Element {
        Element { kind: kind.into(), name: name.into(), value: String::new(), x, y, w: 40.0, h: 20.0, active, flags: vec![], window: "W".into() }
    }

    #[test]
    fn controls_are_numbered_in_reading_order() {
        let els = vec![el("按鈕", "下", 10.0, 300.0, true), el("連結", "Profile", 20.0, 100.0, true), el("文字", "標題", 0.0, 50.0, false), el("連結", "Profile", 20.0, 100.0, true)];
        let (numbered, lines, spots) = number(&els, false);
        // Reading order, duplicates once, plain text left out when there is a picture.
        assert_eq!(lines, vec!["[1] 連結「Profile」", "[2] 按鈕「下」"]);
        assert_eq!(numbered.len(), 2);
        assert_eq!(spots.get("1").unwrap(), ((40.0, 110.0), "連結「Profile」".to_string()));
        assert_eq!(spots.get("e2").unwrap().0, (30.0, 310.0));
        assert!(spots.get("3").is_err() && spots.get("ok").is_err());
        // Without a picture the text is listed too (not numbered).
        let (_, lines, _) = number(&els, true);
        assert_eq!(lines, vec!["    文字「標題」", "[1] 連結「Profile」", "[2] 按鈕「下」"]);
    }

    #[test]
    fn numbers_are_drawn_on_the_picture() {
        let mut img = image::RgbImage::from_pixel(200, 100, image::Rgb([255, 255, 255]));
        let els = vec![el("按鈕", "OK", 100.0, 60.0, true)];
        let (numbered, _, _) = number(&els, false);
        // Screen twice the picture's size: the box lands at (50, 30)–(70, 40).
        draw(&mut img, &numbered, (2.0, 2.0));
        assert_eq!(img.get_pixel(50, 35).0, COLOURS[1]);
        assert_eq!(img.get_pixel(60, 30).0, COLOURS[1]);
        // The label sits above the box with white digits in it.
        let label_has_white = (16..30).any(|y| (50..62).any(|x| img.get_pixel(x, y).0 == [255, 255, 255]));
        let label_has_colour = (16..30).any(|y| (50..62).any(|x| img.get_pixel(x, y).0 == COLOURS[1]));
        assert!(label_has_white && label_has_colour);
    }

    #[test]
    fn a_still_screen_has_the_same_fingerprint() {
        let a = vec![el("按鈕", "OK", 1.0, 2.0, true)];
        let mut b = a.clone();
        assert_eq!(fingerprint(&a), fingerprint(&b));
        b[0].name = "Cancel".into();
        assert_ne!(fingerprint(&a), fingerprint(&b));
        let img = image::RgbImage::from_pixel(64, 40, image::Rgb([10, 20, 30]));
        let mut img2 = img.clone();
        img2.put_pixel(0, 0, image::Rgb([12, 20, 30]));
        assert_eq!(picture_fingerprint(&img), picture_fingerprint(&img2));
    }
}