3 files changed, 72 insertions, 1 deletions

diff --git a/firmware/include/buflib.h b/firmware/include/buflib.h
index c4865b6d9b..3fe8ac1430 100644
--- a/firmware/include/buflib.h
+++ b/firmware/include/buflib.h
@@ -317,6 +317,32 @@ static inline void *buflib_get_data(struct buflib_context *ctx, int handle);
 #endif
 
 /**
+ * \brief Get a pinned pointer to a buflib allocation
+ * \param ctx       Buflib context of the allocation
+ * \param handle    Handle identifying the allocation
+ * \return Pointer to the allocation's memory.
+ *
+ * Functionally equivalent to buflib_pin() followed by buflib_get_data(),
+ * but this call is more efficient and should be preferred over separate
+ * calls.
+ *
+ * To unpin the data, call buflib_put_data_pinned() and pass the pointer
+ * returned by this function.
+ */
+static inline void *buflib_get_data_pinned(struct buflib_context *ctx, int handle);
+
+/**
+ * \brief Release a pinned pointer to a buflib allocation
+ * \param ctx       Buflib context of the allocation
+ * \param data      Pointer returned by buflib_get_data()
+ *
+ * Decrements the pin count, allowing the buffer to be moved once the
+ * pin count drops to zero. This is more efficient than buflib_unpin()
+ * and should be preferred when you have a pointer to the buflib data.
+ */
+static inline void buflib_put_data_pinned(struct buflib_context *ctx, void *data);
+
+/**
  * \brief Shift allocations up to free space at the start of the pool
  * \param ctx       Context to operate on
  * \param size      Indicates number of bytes to free up, or 0 to free
diff --git a/firmware/include/buflib_malloc.h b/firmware/include/buflib_malloc.h
index 32c837e7b7..a17c75c29a 100644
--- a/firmware/include/buflib_malloc.h
+++ b/firmware/include/buflib_malloc.h
@@ -50,4 +50,16 @@ static inline void *buflib_get_data(struct buflib_context *ctx, int handle)
 }
 #endif
 
+static inline void *buflib_get_data_pinned(struct buflib_context *ctx, int handle)
+{
+    buflib_pin(ctx, handle);
+    return buflib_get_data(ctx, handle);
+}
+
+void _buflib_malloc_put_data_pinned(struct buflib_context *ctx, void *data);
+static inline void buflib_put_data_pinned(struct buflib_context *ctx, void *data)
+{
+    _buflib_malloc_put_data_pinned(ctx, data);
+}
+
 #endif /* _BUFLIB_MALLOC_H_ */
diff --git a/firmware/include/buflib_mempool.h b/firmware/include/buflib_mempool.h
index 4b01b629c3..448e40963a 100644
--- a/firmware/include/buflib_mempool.h
+++ b/firmware/include/buflib_mempool.h
@@ -30,6 +30,17 @@
 # error "include buflib.h instead"
 #endif
 
+#include "system.h"
+
+/* Indices used to access block fields as block[BUFLIB_IDX_XXX] */
+enum {
+    BUFLIB_IDX_LEN,     /* length of the block, must come first */
+    BUFLIB_IDX_HANDLE,  /* pointer to entry in the handle table */
+    BUFLIB_IDX_OPS,     /* pointer to an ops struct */
+    BUFLIB_IDX_PIN,     /* pin count */
+    BUFLIB_NUM_FIELDS,
+};
+
 union buflib_data
 {
     intptr_t val;                 /* length of the block in n*sizeof(union buflib_data).
@@ -52,7 +63,7 @@ struct buflib_context
     bool compact;
 };
 
-#define BUFLIB_ALLOC_OVERHEAD (4 * sizeof(union buflib_data))
+#define BUFLIB_ALLOC_OVERHEAD (BUFLIB_NUM_FIELDS * sizeof(union buflib_data))
 
 #ifndef BUFLIB_DEBUG_GET_DATA
 static inline void *buflib_get_data(struct buflib_context *ctx, int handle)
@@ -61,4 +72,26 @@ static inline void *buflib_get_data(struct buflib_context *ctx, int handle)
 }
 #endif
 
+static inline union buflib_data *_buflib_get_block_header(void *data)
+{
+    union buflib_data *bd = ALIGN_DOWN(data, sizeof(*bd));
+    return bd - BUFLIB_NUM_FIELDS;
+}
+
+static inline void *buflib_get_data_pinned(struct buflib_context *ctx, int handle)
+{
+    void *data = buflib_get_data(ctx, handle);
+    union buflib_data *bd = _buflib_get_block_header(data);
+
+    bd[BUFLIB_IDX_PIN].pincount++;
+    return data;
+}
+
+static inline void buflib_put_data_pinned(struct buflib_context *ctx, void *data)
+{
+    (void)ctx;
+    union buflib_data *bd = _buflib_get_block_header(data);
+    bd[BUFLIB_IDX_PIN].pincount--;
+}
+
 #endif /* _BUFLIB_MEMPOOL_H_ */