diff --git a/speculos/mcu/display.py b/speculos/mcu/display.py index a2749c27..666573c2 100644 --- a/speculos/mcu/display.py +++ b/speculos/mcu/display.py @@ -2,6 +2,7 @@ import io from abc import ABC, abstractmethod +from functools import cache from PIL import Image from socket import socket from typing import Any, Dict, IO, List, Optional, Tuple, Union @@ -54,35 +55,6 @@ def can_read(self, screen: DisplayNotifier) -> None: } -def _screenshot_to_iobytes_value(screen_size, data): - image = Image.frombytes("RGB", screen_size, data) - iobytes = io.BytesIO() - image.save(iobytes, format="PNG") - return iobytes.getvalue() - - -class Screenshot: - def __init__(self, screen_size: Tuple[int, int]): - self.pixels: Dict[Tuple[int, int], int] = {} - self.width, self.height = screen_size - for y in range(0, self.height): - for x in range(0, self.width): - self.pixels[(x, y)] = 0x000000 - - def update(self, pixels: Dict[Tuple[int, int], int]) -> None: - # Don't call update, replace the object instead - self.pixels = {**self.pixels, **pixels} - - def get_image(self) -> Tuple[Tuple[int, int], bytes]: - # Get the pixels object once, as it may be replaced during the loop. - data = bytearray(self.width * self.height * 3) - for y in range(0, self.height): - for x in range(0, self.width): - pos = 3 * (y * self.width + x) - data[pos:pos + 3] = self.pixels[(x, y)].to_bytes(3, "big") - return (self.width, self.height), bytes(data) - - class FrameBuffer: """ A class responsible for managing the graphic screen of the current application. @@ -99,34 +71,67 @@ class FrameBuffer: def __init__(self, model: str): self.pixels: Dict[Tuple[int, int], int] = {} + self.screenshot_pixels: Dict[Tuple[int, int], int] = {} + self.default_color = 0 + self.draw_default_color = False self._public_screenshot_value = b'' self.current_data = b'' self.recreate_public_screenshot = True self.model = model self.current_screen_size = MODELS[model].screen_size - self.screenshot = Screenshot(self.current_screen_size) - # Init published content now, don't wait for the first request - if self.model == "stax": - self.update_public_screenshot() + self._width, self._height = MODELS[model].screen_size - def draw_point(self, x: int, y: int, color: int) -> None: + @cache + def check_color(self, color: int) -> int: # There are only 2 colors on the Nano S and the Nano X but the one # passed in argument isn't always valid. Fix it here. if self.model != 'stax': if color != 0x000000: color = FrameBuffer.COLORS.get(self.model, color) - self.pixels[(x, y)] = color + return color + + def draw_point(self, x: int, y: int, color: int) -> None: + self.pixels[(x, y)] = self.check_color(color) + + def draw_horizontal_line(self, x0: int, y: int, width: int, color: int) -> None: + for x in range(x0, x0 + width): + self.pixels[(x, y)] = self.check_color(color) + + def draw_rect(self, x0: int, y0: int, width: int, height: int, color: int) -> None: + color = self.check_color(color) + + if x0 == 0 and y0 == 0 and width == self._width and height == self._height: + self.default_color = color + self.draw_default_color = True + self.pixels = {} + self.screenshot_pixels = {} + return + + for x in range(x0, x0 + width): + for y in range(y0, y0 + height): + self.pixels[(x, y)] = color + + def _get_image(self) -> bytes: + data = bytearray(self.default_color.to_bytes(3, "big")) * self._width * self._height + for (x, y), color in self.screenshot_pixels.items(): + pos = 3 * (y * self._width + x) + data[pos:pos + 3] = color.to_bytes(3, "big") + return bytes(data) + + def _get_screenshot_iobytes_value(self) -> bytes: + # Get the pixels object once, as it may be replaced during the loop. + data = self._get_image() - def screenshot_update_pixels(self): - # Update the screenshot object with our current pixels content - self.screenshot.update(self.pixels) + image = Image.frombytes("RGB", self.current_screen_size, data) + iobytes = io.BytesIO() + image.save(iobytes, format="PNG") + return iobytes.getvalue() def take_screenshot(self) -> Tuple[Tuple[int, int], bytes]: - self.current_screen_size, self.current_data = self.screenshot.get_image() - return self.current_screen_size, self.current_data + return self.current_screen_size, self._get_image() def update_screenshot(self) -> None: - self.screenshot.update(self.pixels) + self.screenshot_pixels = {**self.screenshot_pixels, **self.pixels} def update_public_screenshot(self) -> None: # Stax only @@ -140,7 +145,7 @@ def public_screenshot_value(self) -> bytes: # and not necessary if no one tries to read the value if self.recreate_public_screenshot: self.recreate_public_screenshot = False - self._public_screenshot_value = _screenshot_to_iobytes_value(self.current_screen_size, self.current_data) + self._public_screenshot_value = self._get_screenshot_iobytes_value() return self._public_screenshot_value @@ -152,8 +157,7 @@ def get_public_screenshot(self) -> bytes: return self.public_screenshot_value # On nano we have no knowledge of screen refreshes so we can't be scarce on publishes # So we publish the raw current content every time. It's ok as take_screenshot is fast on Nano - screen_size, data = self.take_screenshot() - return _screenshot_to_iobytes_value(screen_size, data) + return self._get_screenshot_iobytes_value() # Should be declared as an `@abstractmethod` (and also `class FrameBuffer(ABC):`), but in this # case multiple inheritance in `screen.PaintWidget(FrameBuffer, QWidget)` will break, as both diff --git a/speculos/mcu/headless.py b/speculos/mcu/headless.py index f0df659b..7a45d53e 100644 --- a/speculos/mcu/headless.py +++ b/speculos/mcu/headless.py @@ -53,15 +53,16 @@ def update(self, _1: Optional[int] = None, _2: Optional[int] = None, _3: Optional[int] = None) -> bool: - if self.pixels: + if self.pixels or self.draw_default_color: self._redraw() self.pixels = {} + self.draw_default_color = False return True return False def _redraw(self) -> None: if self.vnc: - self.vnc.redraw(self.pixels) + self.vnc.redraw(self.pixels, self.default_color) self.update_screenshot() diff --git a/speculos/mcu/nbgl.py b/speculos/mcu/nbgl.py index 944a2ec4..079bd2bd 100644 --- a/speculos/mcu/nbgl.py +++ b/speculos/mcu/nbgl.py @@ -3,6 +3,7 @@ import sys from construct import Struct, Int8ul, Int16ul from enum import IntEnum +from functools import cache from typing import Tuple from .display import FrameBuffer, GraphicLibrary @@ -30,6 +31,7 @@ class NbglColor(IntEnum): class NBGL(GraphicLibrary): @staticmethod + @cache def to_screen_color(color: int, bpp: int) -> int: color_table = { 1: 0xFFFFFF, @@ -57,9 +59,7 @@ def __assert_area(self, area) -> None: def hal_draw_rect(self, data: bytes) -> None: area = nbgl_area_t.parse(data) self.__assert_area(area) - for x in range(area.x0, area.x0+area.width): - for y in range(area.y0, area.y0+area.height): - self.fb.draw_point(x, y, NBGL.to_screen_color(area.color, 2)) + self.fb.draw_rect(area.x0, area.y0, area.width, area.height, NBGL.to_screen_color(area.color, 2)) def refresh(self, data: bytes) -> bool: area = nbgl_area_t.parse(data) @@ -74,19 +74,21 @@ def hal_draw_line(self, data: bytes) -> None: back_color = NBGL.to_screen_color(area.color, 2) front_color = NBGL.to_screen_color(color, 2) - for x in range(area.x0, area.x0+area.width): - for y in range(area.y0, area.y0+area.height): - if (mask >> (y-area.y0)) & 0x1: - self.fb.draw_point(x, y, front_color) - else: - self.fb.draw_point(x, y, back_color) + + for y in range(area.y0, area.y0+area.height): + if (mask >> (y-area.y0)) & 0x1: + self.fb.draw_horizontal_line(area.x0, y, area.width, front_color) + else: + self.fb.draw_horizontal_line(area.x0, y, area.width, back_color) @staticmethod + @cache def get_color_from_color_map(color, color_map, bpp): # #define GET_COLOR_MAP(__map__,__col__) ((__map__>>(__col__*2))&0x3) return NBGL.to_screen_color((color_map >> (color*2)) & 0x3, bpp) @staticmethod + @cache def get_4bpp_color_from_color_index(index, front_color, back_color): COLOR_MAPS_4BPP = { # Manually hardcoced color maps diff --git a/speculos/mcu/screen.py b/speculos/mcu/screen.py index 9d00a69b..20415e18 100644 --- a/speculos/mcu/screen.py +++ b/speculos/mcu/screen.py @@ -30,7 +30,7 @@ def __init__(self, parent, model: str, pixel_size: int, vnc: Optional[VNC] = Non self.vnc = vnc def paintEvent(self, event: QEvent): - if self.pixels: + if self.pixels or self.draw_default_color: pixmap = QPixmap(self.size() / self.pixel_size) pixmap.fill(Qt.white) painter = QPainter(pixmap) @@ -38,6 +38,7 @@ def paintEvent(self, event: QEvent): self._redraw(painter) self.mPixmap = pixmap self.pixels = {} + self.draw_default_color = False qp = QPainter(self) copied_pixmap = self.mPixmap @@ -60,12 +61,15 @@ def update(self, # type: ignore[override] return self.pixels != {} def _redraw(self, qp): + if self.draw_default_color: + qp.fillRect(0, 0, self._width, self._height, QColor.fromRgb(self.default_color)) + for (x, y), color in self.pixels.items(): qp.setPen(QColor.fromRgb(color)) qp.drawPoint(x, y) if self.vnc is not None: - self.vnc.redraw(self.pixels) + self.vnc.redraw(self.pixels, self.default_color) self.update_screenshot() diff --git a/speculos/mcu/seproxyhal.py b/speculos/mcu/seproxyhal.py index edca13ce..9f64af0a 100644 --- a/speculos/mcu/seproxyhal.py +++ b/speculos/mcu/seproxyhal.py @@ -322,11 +322,8 @@ def can_read(self, screen: DisplayNotifier): if self.refreshed: self.refreshed = False - # Run the OCR + # Update the screenshot, we'll upload its associated events shortly screen.display.gl.update_screenshot() - screen.display.gl.take_screenshot() - - # Publish the new screenshot, we'll upload its associated events shortly screen.display.gl.update_public_screenshot() if screen.display.model != "stax" and screen.display.screen_update(): @@ -423,7 +420,7 @@ def can_read(self, screen: DisplayNotifier): screen.display.gl.refresh(data) # Stax only # We have refreshed the screen, remember it for the next time we have SephTag.GENERAL_STATUS - # then we'll perform a new OCR and make public the resulting screenshot / OCR analysis + # then we'll perform a screen update and make public the resulting screenshot self.refreshed = True elif tag == SephTag.NBGL_DRAW_LINE: diff --git a/speculos/mcu/vnc.py b/speculos/mcu/vnc.py index 7c4af075..472fbf46 100644 --- a/speculos/mcu/vnc.py +++ b/speculos/mcu/vnc.py @@ -22,13 +22,13 @@ def __init__(self, verbose: bool = False): self.logger = logging.getLogger("vnc") - width, height = screen_size + self.width, self.height = screen_size path = os.path.dirname(os.path.realpath(__file__)) server = os.path.join(path, '../resources/vnc_server') cmd = [server] # custom options - cmd += ['-s', f'{width}x{height}'] + cmd += ['-s', f'{self.width}x{self.height}'] if verbose: cmd += ['-v'] @@ -48,23 +48,25 @@ def file(self) -> IO[bytes]: assert self.subprocess.stdout is not None return self.subprocess.stdout - def redraw(self, pixels): + def redraw(self, pixels, default_color): '''The framebuffer was updated, forward everything to the VNC server.''' # int.to_bytes() is super slow, hence the manual encoding buf = bytearray(len(pixels) * 9) i = 0 - for (x, y), color in pixels.items(): - buf[i + 0] = y & 0xff - buf[i + 1] = (y >> 8) & 0xff - buf[i + 2] = x & 0xff - buf[i + 3] = (x >> 8) & 0xff - buf[i + 4] = color & 0xff - buf[i + 5] = (color >> 8) & 0xff - buf[i + 6] = (color >> 16) & 0xff - buf[i + 7] = (color >> 24) & 0xff - buf[i + 8] = 0x0a - i += 9 + for x in range(0, self._width): + for y in range(0, self._height): + color = pixels.get((x, y), default_color) + buf[i + 0] = y & 0xff + buf[i + 1] = (y >> 8) & 0xff + buf[i + 2] = x & 0xff + buf[i + 3] = (x >> 8) & 0xff + buf[i + 4] = color & 0xff + buf[i + 5] = (color >> 8) & 0xff + buf[i + 6] = (color >> 16) & 0xff + buf[i + 7] = (color >> 24) & 0xff + buf[i + 8] = 0x0a + i += 9 self.subprocess.stdin.write(buf) self.subprocess.stdin.flush()