Add methods to get locked buffer from Surface.

Bug: 307535208
Test: atest libnativewindow_rs-internal_test
Change-Id: Ib9a824a72d1d415dddd82e477d55200765d0b9e6

3 files changed, 91 insertions, 6 deletions

diff --git a/libs/nativewindow/rust/src/lib.rs b/libs/nativewindow/rust/src/lib.rs
index 931c311e65..f19b908074 100644
--- a/libs/nativewindow/rust/src/lib.rs
+++ b/libs/nativewindow/rust/src/lib.rs
@@ -19,10 +19,9 @@ extern crate nativewindow_bindgen as ffi;
 mod handle;
 mod surface;
 
-pub use handle::NativeHandle;
-pub use surface::Surface;
-
 pub use ffi::{AHardwareBuffer_Format, AHardwareBuffer_UsageFlags};
+pub use handle::NativeHandle;
+pub use surface::{buffer::Buffer, Surface};
 
 use binder::{
     binder_impl::{BorrowedParcel, UnstructuredParcelable},
diff --git a/libs/nativewindow/rust/src/surface.rs b/libs/nativewindow/rust/src/surface.rs
index 9eddfcd497..ed52471aa1 100644
--- a/libs/nativewindow/rust/src/surface.rs
+++ b/libs/nativewindow/rust/src/surface.rs
@@ -14,6 +14,8 @@
 
 //! Rust wrapper for `ANativeWindow` and related types.
 
+pub(crate) mod buffer;
+
 use binder::{
     binder_impl::{BorrowedParcel, UnstructuredParcelable},
     impl_deserialize_for_unstructured_parcelable, impl_serialize_for_unstructured_parcelable,
@@ -21,17 +23,18 @@ use binder::{
     StatusCode,
 };
 use bitflags::bitflags;
+use buffer::Buffer;
 use nativewindow_bindgen::{
     ADataSpace, AHardwareBuffer_Format, ANativeWindow, ANativeWindow_acquire,
     ANativeWindow_getBuffersDataSpace, ANativeWindow_getBuffersDefaultDataSpace,
-    ANativeWindow_getFormat, ANativeWindow_getHeight, ANativeWindow_getWidth,
+    ANativeWindow_getFormat, ANativeWindow_getHeight, ANativeWindow_getWidth, ANativeWindow_lock,
     ANativeWindow_readFromParcel, ANativeWindow_release, ANativeWindow_setBuffersDataSpace,
     ANativeWindow_setBuffersGeometry, ANativeWindow_setBuffersTransform,
-    ANativeWindow_writeToParcel,
+    ANativeWindow_unlockAndPost, ANativeWindow_writeToParcel, ARect,
 };
 use std::error::Error;
 use std::fmt::{self, Debug, Display, Formatter};
-use std::ptr::{null_mut, NonNull};
+use std::ptr::{self, null_mut, NonNull};
 
 /// Wrapper around an opaque C `ANativeWindow`.
 #[derive(PartialEq, Eq)]
@@ -153,6 +156,43 @@ impl Surface {
             Ok(ADataSpace(data_space))
         }
     }
+
+    /// Locks the window's next drawing surface for writing, and returns it.
+    pub fn lock(&mut self, bounds: Option<&mut ARect>) -> Result<Buffer, ErrorCode> {
+        let mut buffer = buffer::EMPTY;
+        // SAFETY: The ANativeWindow pointer we pass is guaranteed to be non-null and valid because
+        // it must have been allocated by `ANativeWindow_allocate` or `ANativeWindow_readFromParcel`
+        // and we have not yet released it. The other pointers must be valid because the come from
+        // references, and aren't retained after the function returns.
+        let status = unsafe {
+            ANativeWindow_lock(
+                self.0.as_ptr(),
+                &mut buffer,
+                bounds.map(ptr::from_mut).unwrap_or(null_mut()),
+            )
+        };
+        if status != 0 {
+            return Err(ErrorCode(status));
+        }
+
+        Ok(Buffer::new(buffer, self))
+    }
+
+    /// Unlocks the window's drawing surface which was previously locked, posting the new buffer to
+    /// the display.
+    ///
+    /// This shouldn't be called directly but via the [`Buffer`], hence is not public here.
+    fn unlock_and_post(&mut self) -> Result<(), ErrorCode> {
+        // SAFETY: The ANativeWindow pointer we pass is guaranteed to be non-null and valid because
+        // it must have been allocated by `ANativeWindow_allocate` or `ANativeWindow_readFromParcel`
+        // and we have not yet released it.
+        let status = unsafe { ANativeWindow_unlockAndPost(self.0.as_ptr()) };
+        if status == 0 {
+            Ok(())
+        } else {
+            Err(ErrorCode(status))
+        }
+    }
 }
 
 impl Drop for Surface {
diff --git a/libs/nativewindow/rust/src/surface/buffer.rs b/libs/nativewindow/rust/src/surface/buffer.rs
new file mode 100644
index 0000000000..91f5de5343
--- /dev/null
+++ b/libs/nativewindow/rust/src/surface/buffer.rs
@@ -0,0 +1,46 @@
+// Copyright (C) 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::{ErrorCode, Surface};
+use nativewindow_bindgen::ANativeWindow_Buffer;
+use std::ptr::null_mut;
+
+/// An empty `ANativeWindow_Buffer`.
+pub const EMPTY: ANativeWindow_Buffer = ANativeWindow_Buffer {
+    width: 0,
+    height: 0,
+    stride: 0,
+    format: 0,
+    bits: null_mut(),
+    reserved: [0; 6],
+};
+
+/// Rust wrapper for `ANativeWindow_Buffer`, representing a locked buffer from a [`Surface`].
+pub struct Buffer<'a> {
+    /// The wrapped `ANativeWindow_Buffer`.
+    pub buffer: ANativeWindow_Buffer,
+    surface: &'a mut Surface,
+}
+
+impl<'a> Buffer<'a> {
+    pub(crate) fn new(buffer: ANativeWindow_Buffer, surface: &'a mut Surface) -> Self {
+        Self { buffer, surface }
+    }
+
+    /// Unlocks the window's drawing surface which was previously locked to create this buffer,
+    /// posting the buffer to the display.
+    pub fn unlock_and_post(self) -> Result<(), ErrorCode> {
+        self.surface.unlock_and_post()
+    }
+}