ai-app/iris/src/android/view.rs

use crate::prelude::*;
use crate::task::RequestRedraw;
use android_view::{
    CallbackCtx, Context, InputConnection, KeyEvent, MotionEvent, Rect, View, ViewPeer,
    jni::JNIEnv,
    ndk::event::{Keycode, MotionAction},
};
// `marker::Sized` explicitly: `crate::prelude::*` below also brings in the
// `Sized` *widget* (`widget::position::sized::Sized`), and an unqualified
// glob import shadows the language prelude -- `default/mod.rs` has the same
// explicit import for the same reason.
use std::{
    cell::RefCell,
    marker::{PhantomData, Sized},
    rc::Rc,
    sync::Arc,
    time::Instant,
};

use super::{
    insets::{Insets, Shared},
    render::{AndroidRedrawHandle, AndroidRenderer},
};

/// The android-view analogue of `default::DefaultUiState`. `renderer` is an
/// `Option` because a `SurfaceView`'s surface does not outlive backgrounding
/// the way a winit `Window` does -- `surfaceDestroyed`/`surfaceCreated` can
/// happen any number of times over the life of one `IrisViewPeer`.
pub struct AndroidUiState {
    pub root: Option<StrongWidget>,
    pub renderer: Option<AndroidRenderer>,
    pub focus: Option<WeakWidget<TextEdit>>,
    pub cursor: CursorState,
    pub last_click: Instant,
    /// The IME preedit's previous length, in `char`s -- the same
    /// re-send-the-whole-composition bookkeeping `default::DefaultUiState`
    /// keeps for winit's `Ime::Preedit`, since android-view's
    /// `setComposingText` has the identical shape (see `android/ime.rs`).
    pub compose_len: usize,
    /// Set by `attr::FocusHost::focus_gained` when a `TextEdit` is focused;
    /// consumed by the touch handler after the sensor pass finishes, since
    /// showing the keyboard is a JNI call and `focus_gained` runs deep
    /// inside the platform-agnostic sensor dispatch with no `CallbackCtx`
    /// in reach.
    pub pending_show_keyboard: bool,
    /// Window insets, filled in from outside the normal `ViewPeer` callback
    /// path -- see `android/insets.rs` for why they need a registry of
    /// their own.
    shared: Rc<RefCell<Shared>>,
}

impl AndroidUiState {
    fn new(shared: Rc<RefCell<Shared>>) -> Self {
        Self {
            root: None,
            renderer: None,
            focus: None,
            cursor: Default::default(),
            last_click: Instant::now(),
            compose_len: 0,
            pending_show_keyboard: false,
            shared,
        }
    }

    pub fn insets(&self) -> Insets {
        self.shared.borrow().insets
    }
}

impl HasRoot for AndroidUiState {
    fn set_root(&mut self, root: StrongWidget) {
        self.root = Some(root);
    }
}

pub trait HasAndroidUiState: Sized + 'static {
    fn android_state(&self) -> &AndroidUiState;
    fn android_state_mut(&mut self) -> &mut AndroidUiState;
}

pub trait AndroidAppState: HasAndroidUiState {
    fn new(ui_state: AndroidUiState, rsc: &mut AndroidRsc<Self>) -> Self;
    /// The system back gesture/button. `true` means handled -- nothing
    /// further happens; `false` lets the activity finish as it would with
    /// no view at all. The default declines, since most screens have
    /// nothing to intercept it for.
    #[allow(unused_variables)]
    fn back_pressed(&mut self, rsc: &mut AndroidRsc<Self>, render: &mut UiRenderState) -> bool {
        false
    }
}

/// The android-view analogue of `default::DefaultRsc` -- identical in
/// substance, since none of `UiRsc`/`HasEvents`/`HasTasks`/`HasWidgetState`
/// mention winit. Kept as a separate type rather than shared code because
/// the two backends' `ViewPeer`/`ApplicationHandler` entry points hold
/// their harness state differently (see RUST.md's I2).
pub struct AndroidRsc<State: 'static> {
    pub ui: UiData,
    pub events: EventManager<Self>,
    pub tasks: Tasks<Self>,
    pub state: WidgetState,
    _state: PhantomData<State>,
}

impl<State> AndroidRsc<State> {
    pub fn create_state<T: 'static>(&mut self, id: impl IdLike, data: T) -> WeakState<T> {
        self.state.add(id.id(), data)
    }
}

impl<State> UiRsc for AndroidRsc<State> {
    fn ui(&self) -> &UiData {
        &self.ui
    }
    fn ui_mut(&mut self) -> &mut UiData {
        &mut self.ui
    }
    fn on_draw(&mut self, active: &ActiveData) {
        self.events.draw(active);
    }
    fn on_undraw(&mut self, active: &ActiveData) {
        self.events.undraw(active);
    }
    fn on_remove(&mut self, id: WidgetId) {
        self.events.remove(id);
        self.state.remove(id);
    }
}

impl<State: 'static> HasState for AndroidRsc<State> {
    type State = State;
}

impl<State: 'static> HasEvents for AndroidRsc<State> {
    fn events(&self) -> &EventManager<Self> {
        &self.events
    }
    fn events_mut(&mut self) -> &mut EventManager<Self> {
        &mut self.events
    }
}

impl<State: 'static> HasTasks for AndroidRsc<State> {
    fn tasks_mut(&mut self) -> &mut Tasks<Self> {
        &mut self.tasks
    }
}

impl<State: 'static> HasWidgetState for AndroidRsc<State> {
    fn widget_state(&self) -> &WidgetState {
        &self.state
    }
    fn widget_state_mut(&mut self) -> &mut WidgetState {
        &mut self.state
    }
}

/// The `ViewPeer` android-view dispatches every callback to. One per
/// `RustView` instance; `new_peer` (below) builds it and hands the id to
/// Java the same way android-view's own demo does.
pub struct IrisViewPeer<State: AndroidAppState> {
    pub(super) rsc: AndroidRsc<State>,
    pub(super) render: UiRenderState,
    pub(super) state: State,
    task_recv: TaskMsgReceiver<AndroidRsc<State>>,
}

impl<State: 'static, I: RscIdx<AndroidRsc<State>>> std::ops::Index<I> for AndroidRsc<State> {
    type Output = I::Output;

    fn index(&self, index: I) -> &Self::Output {
        index.get(self)
    }
}

impl<State: 'static, I: RscIdx<AndroidRsc<State>>> std::ops::IndexMut<I> for AndroidRsc<State> {
    fn index_mut(&mut self, index: I) -> &mut Self::Output {
        index.get_mut(self)
    }
}

impl<State: AndroidAppState> IrisViewPeer<State> {
    fn drain_tasks(&mut self) {
        while let Ok(update) = self.task_recv.try_recv() {
            update(&mut self.state, &mut self.rsc);
        }
    }

    /// Common tail for every callback that might have changed the cursor,
    /// the text focus, or the widget tree: run the sensors that touch
    /// input feeds, then ask for a frame if the result needs drawing.
    /// Mirrors `default::DefaultApp::window_event`'s tail, split across
    /// android-view's several entry points instead of winit's one.
    pub(super) fn after_input(&mut self, ctx: &mut CallbackCtx) {
        let window_size = self.window_size();
        let ui_state = self.state.android_state_mut();
        let cursor = ui_state.cursor.clone();
        let old_focus = ui_state.focus;
        self.render
            .run_sensors(&mut self.rsc, &mut self.state, cursor, window_size);

        let ui_state = self.state.android_state_mut();
        if old_focus != ui_state.focus
            && let Some(old) = old_focus
        {
            old.edit(&mut self.rsc).deselect();
        }
        if std::mem::take(&mut ui_state.pending_show_keyboard) {
            show_soft_input(&mut ctx.env, &ctx.view);
        }

        let ui_state = self.state.android_state_mut();
        ui_state.cursor.end_frame();
        if self.render.needs_redraw(&ui_state.root, self.rsc.widgets()) {
            ctx.view.post_frame_callback(&mut ctx.env);
        }
    }

    fn window_size(&self) -> Vec2 {
        let ui_state = self.state.android_state();
        match &ui_state.renderer {
            Some(r) => r.size(),
            None => Vec2::ZERO,
        }
    }

    /// The `log::debug!` calls here are a live diagnostic for a still-open
    /// finding (RUST.md's I2): layout runs and reports the right pixel
    /// region for the root (confirmed via `window_region`, logged below),
    /// and the clear colour reaches the screen (confirmed by swapping it to
    /// magenta and screenshotting), but no primitive ever appears on top of
    /// it -- on both the Vulkan/SwiftShader and GLES/virgl backends. Leave
    /// these in until that is root-caused; removing them loses the exact
    /// evidence a `logcat` capture needs to reproduce the state.
    fn render(&mut self) {
        let ui_state = self.state.android_state();
        if ui_state.renderer.is_none() {
            return;
        }
        log::debug!(
            "render(): root={:?} widgets={} active={} root_px={:?} out_size={:?}",
            ui_state.root.is_some(),
            self.rsc.widgets().len(),
            self.render.active_widgets(),
            ui_state
                .root
                .as_ref()
                .and_then(|r| self.render.window_region(r, &self.rsc)),
            self.window_size(),
        );
        let ui_state = self.state.android_state_mut();
        self.render.update(&ui_state.root, &mut self.rsc);
        let ui_state = self.state.android_state_mut();
        let Some(renderer) = &mut ui_state.renderer else {
            return;
        };
        renderer.update(&mut self.rsc.ui, &mut self.render);
        renderer.draw();
        let ui_state = self.state.android_state();
        log::debug!(
            "render(): after update active={} root_px={:?}",
            self.render.active_widgets(),
            ui_state
                .root
                .as_ref()
                .and_then(|r| self.render.window_region(r, &self.rsc)),
        );
    }
}

fn show_soft_input<'local>(env: &mut JNIEnv<'local>, view: &View<'local>) {
    let imm = view.input_method_manager(env);
    imm.show_soft_input(env, view, 0);
}

impl<State: AndroidAppState> ViewPeer for IrisViewPeer<State> {
    fn on_key_down<'local>(
        &mut self,
        ctx: &mut CallbackCtx<'local>,
        key_code: Keycode,
        event: &KeyEvent<'local>,
    ) -> bool {
        self.drain_tasks();
        // With no `OnBackPressedCallback` registered on the Java side, the
        // system still delivers the back gesture as a synthetic
        // `KEYCODE_BACK` through this same path -- the legacy behaviour
        // every view-based app gets by default, and enough for "the back
        // gesture as an event" without a second JNI registry. See
        // `android/insets.rs`'s doc comment for why insets could not take
        // the same shortcut.
        if key_code == Keycode::Back {
            let handled = self.state.back_pressed(&mut self.rsc, &mut self.render);
            if handled {
                self.after_input(ctx);
            }
            return handled;
        }
        let handled = super::input::on_key(
            &mut self.rsc,
            &mut self.state,
            &mut ctx.env,
            key_code,
            event,
        );
        if handled {
            self.after_input(ctx);
        }
        handled
    }

    fn on_touch_event<'local>(
        &mut self,
        ctx: &mut CallbackCtx<'local>,
        event: &MotionEvent<'local>,
    ) -> bool {
        self.drain_tasks();
        let action = event.action_masked(&mut ctx.env);
        let x = event.x(&mut ctx.env);
        let y = event.y(&mut ctx.env);
        let ui_state = self.state.android_state_mut();
        match action {
            MotionAction::Down => {
                ui_state.cursor.pos = vec2(x, y);
                ui_state.cursor.exists = true;
                ui_state.cursor.buttons.left.update(true);
            }
            MotionAction::Move => {
                ui_state.cursor.pos = vec2(x, y);
            }
            MotionAction::Up | MotionAction::Cancel => {
                ui_state.cursor.pos = vec2(x, y);
                ui_state.cursor.buttons.left.update(false);
            }
            _ => return false,
        }
        self.after_input(ctx);
        true
    }

    fn on_focus_changed<'local>(
        &mut self,
        ctx: &mut CallbackCtx<'local>,
        gain_focus: bool,
        _direction: i32,
        _previously_focused_rect: Option<&Rect<'local>>,
    ) {
        self.drain_tasks();
        if !gain_focus {
            let ui_state = self.state.android_state_mut();
            if let Some(focus) = ui_state.focus.take() {
                focus.edit(&mut self.rsc).deselect();
            }
        }
        self.after_input(ctx);
    }

    fn on_attached_to_window(&mut self, _ctx: &mut CallbackCtx) {
        self.drain_tasks();
    }

    fn surface_changed<'local>(
        &mut self,
        ctx: &mut CallbackCtx<'local>,
        holder: &android_view::SurfaceHolder<'local>,
        _format: i32,
        width: i32,
        height: i32,
    ) {
        self.drain_tasks();
        let window = holder.surface(&mut ctx.env).to_native_window(&mut ctx.env);
        // The layout engine's own notion of the canvas size is separate
        // from the wgpu surface's -- winit's backend sets it from
        // `WindowEvent::Resized`, and there is no equivalent automatic
        // trigger here, so this is the one place android-view's surface
        // size has to be told to `UiRenderState` too. Missing this drew
        // nothing but the clear colour: the widget tree laid out against
        // whatever size `UiRenderState::new` starts at instead of the
        // surface's real one.
        self.render.resize((width as u32, height as u32));
        // Drop the old renderer (and the surface it owns) before building
        // one from the new window -- see `AndroidRenderer`'s doc comment.
        let ui_state = self.state.android_state_mut();
        ui_state.renderer = None;
        ui_state.renderer = Some(AndroidRenderer::new(window, width as u32, height as u32));
        self.render();
    }

    fn surface_destroyed<'local>(
        &mut self,
        _ctx: &mut CallbackCtx<'local>,
        _holder: &android_view::SurfaceHolder<'local>,
    ) {
        self.state.android_state_mut().renderer = None;
    }

    fn do_frame(&mut self, _ctx: &mut CallbackCtx, _frame_time_nanos: i64) {
        self.drain_tasks();
        self.render();
    }

    fn as_input_connection(&mut self) -> Option<&mut dyn InputConnection> {
        Some(self)
    }
}

/// Registers `IrisViewPeer<State>`'s native methods and builds one on every
/// `newViewPeer` call from Java. `State`'s app crate wraps this in a
/// concrete `extern "system" fn` (a generic function cannot be handed to
/// `register_view_class`, which wants a plain function pointer) -- see
/// `iris/android-app/src/lib.rs`.
pub fn new_peer<'local, State: AndroidAppState>(
    env: JNIEnv<'local>,
    view: View<'local>,
    _context: Context<'local>,
) -> android_view::jni::sys::jlong {
    let vm = env.get_java_vm().unwrap();
    let global_view = env.new_global_ref(&view.0).unwrap();
    let redraw: Arc<dyn RequestRedraw> = Arc::new(AndroidRedrawHandle::new(vm, global_view));
    let (tasks, task_recv) = Tasks::init(redraw);
    let mut rsc = AndroidRsc {
        ui: Default::default(),
        events: Default::default(),
        tasks,
        state: Default::default(),
        _state: PhantomData,
    };
    let shared = Rc::new(RefCell::new(Shared::default()));
    let ui_state = AndroidUiState::new(shared.clone());
    let state = State::new(ui_state, &mut rsc);
    let peer = IrisViewPeer {
        rsc,
        render: UiRenderState::new(),
        state,
        task_recv,
    };
    let id = android_view::register_view_peer(peer);
    super::insets::register(id, shared);
    id
}