Add crude ws2812b driver.

This driver can run lights! However it is very crude and does not
provide a great API. It does work though! Currently lights up 5 leds of
different colors.

Right now this does synchronous writes, which is not the best. It would
be better to be async, but I think that sholud not be a problem and will
shortly be done.

6 files changed, 128 insertions, 16 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9029fdc..a7c126d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2,5 +2,7 @@
 # CMakeLists in this exact order for cmake to work correctly
 cmake_minimum_required(VERSION 3.5)
 
+list(APPEND EXTRA_COMPONENT_DIRS "${CMAKE_CURRENT_LIST_DIR}/main/drv")
+
 include($ENV{IDF_PATH}/tools/cmake/project.cmake)
 project(ws2812b)
diff --git a/include/drv/ws2812b.h b/include/drv/ws2812b.h
new file mode 100644
index 0000000..9cf35e4
--- /dev/null
+++ b/include/drv/ws2812b.h
@@ -0,0 +1,42 @@
+#pragma once
+#ifndef INCLUDE_DRV_WS2812B_H_
+#define INCLUDE_DRV_WS2812B_H_
+
+#include <stdint.h>
+#include "driver/spi_master.h"
+
+typedef struct {
+  uint8_t r;
+  uint8_t g;
+  uint8_t b;
+} ws2812b_rgb_t;
+
+typedef struct {
+  uint8_t* buf_;       /* Buffer to hold the value to spi. */
+
+  uint32_t n_rgb;      /* Number of rgb values which exist. */
+  ws2812b_rgb_t rgb[]; /* Colors to write. */
+} ws2812b_buffer_t;
+
+struct WS2812B;
+typedef struct WS2812B ws2812b_t;
+
+ws2812b_buffer_t* ws2812b_new_buffer(uint32_t size);
+
+ws2812b_t* ws2812b_init(spi_device_handle_t spi);
+
+esp_err_t ws2812b_write_sync(ws2812b_t* drv, ws2812b_buffer_t* buffer);
+
+static inline void ws2812b_buffer_set_rgb(
+    ws2812b_buffer_t* buf, size_t idx, uint8_t r, uint8_t g, uint8_t b)
+{
+  if (idx >= buf->n_rgb) {
+    return;
+  }
+
+  buf->rgb[idx].r = r;
+  buf->rgb[idx].g = g;
+  buf->rgb[idx].b = b;
+}
+
+#endif /* INCLUDE_DRV_WS2812B_H_ */
diff --git a/main/CMakeLists.txt b/main/CMakeLists.txt
index cd28b28..f7f0237 100644
--- a/main/CMakeLists.txt
+++ b/main/CMakeLists.txt
@@ -1,2 +1,2 @@
 idf_component_register(SRCS "main.c"
-  INCLUDE_DIRS "")
+  INCLUDE_DIRS "../include")
diff --git a/main/drv/CMakeLists.txt b/main/drv/CMakeLists.txt
new file mode 100644
index 0000000..b56f1b0
--- /dev/null
+++ b/main/drv/CMakeLists.txt
@@ -0,0 +1,2 @@
+idf_component_register(SRCS "ws2812b.c"
+  INCLUDE_DIRS "../../include")
diff --git a/main/drv/ws2812b.c b/main/drv/ws2812b.c
new file mode 100644
index 0000000..aaa3c07
--- /dev/null
+++ b/main/drv/ws2812b.c
@@ -0,0 +1,68 @@
+#include <drv/ws2812b.h>
+
+struct WS2812B {
+  spi_device_handle_t spi_handle;
+};
+
+ws2812b_t* ws2812b_init(spi_device_handle_t spi)
+{
+  ws2812b_t* ret = heap_caps_malloc(sizeof(ws2812b_t), MALLOC_CAP_DEFAULT);
+  ret->spi_handle = spi;
+
+  return ret;
+}
+
+/* Compiles a byte into it's padded serial format. */
+#undef BIT
+#define BIT(b, n) (!!((b) & (1 << (n))))
+static inline void byte_compile(uint8_t byte, uint8_t compl [3])
+{
+  compl [0] = 0 | 1 << 7 | BIT(byte, 7) << 6 | 0 << 5 | 1 << 4 |
+              BIT(byte, 6) << 3 | 0 << 2 | 1 << 1 | BIT(byte, 5) << 0;
+  compl [1] = 0 | 0 << 7 | 1 << 6 | BIT(byte, 4) << 5 | 0 << 4 | 1 << 3 |
+              BIT(byte, 3) << 2 | 0 << 1 | 1 << 0;
+  compl [2] = 0 | BIT(byte, 2) << 7 | 0 << 6 | 1 << 5 | BIT(byte, 1) << 4 |
+              0 << 3 | 1 << 2 | BIT(byte, 0) << 1 | 0 << 0;
+}
+
+static inline void compile(ws2812b_buffer_t* buffer)
+{
+  for (size_t i = 0; i < buffer->n_rgb; ++ i) {
+    byte_compile(buffer->rgb[i].g, buffer->buf_ + i * 9 + 0);
+    byte_compile(buffer->rgb[i].r, buffer->buf_ + i * 9 + 3);
+    byte_compile(buffer->rgb[i].b, buffer->buf_ + i * 9 + 6);
+  }
+}
+
+ws2812b_buffer_t* ws2812b_new_buffer(uint32_t n)
+{
+  ws2812b_buffer_t* ret = heap_caps_malloc(
+      sizeof(ws2812b_buffer_t) + sizeof(ws2812b_rgb_t) * n, MALLOC_CAP_DEFAULT);
+  ret->buf_ = heap_caps_malloc(n * 9, MALLOC_CAP_DMA);
+  ret->n_rgb = n;
+  return ret;
+}
+
+esp_err_t ws2812b_write_sync(ws2812b_t* drv, ws2812b_buffer_t* buffer)
+{
+  esp_err_t err;
+
+  compile(buffer);
+  spi_transaction_t t = { 0 };
+  spi_transaction_t* rt;
+
+  t.tx_buffer = buffer->buf_;
+  t.length = buffer->n_rgb * 9 * 8;
+
+  err = spi_device_queue_trans(drv->spi_handle, &t, portMAX_DELAY);
+  if (err != ESP_OK) {
+    return err;
+  }
+
+  err = spi_device_get_trans_result(drv->spi_handle, &rt, portMAX_DELAY);
+  if (err != ESP_OK) {
+    return err;
+  }
+
+  return ESP_OK;
+}
diff --git a/main/main.c b/main/main.c
index b4c744f..cffa6f5 100644
--- a/main/main.c
+++ b/main/main.c
@@ -5,6 +5,8 @@
 #include "esp_spi_flash.h"
 #include "driver/spi_master.h"
 
+#include "drv/ws2812b.h"
+
 #ifdef CONFIG_IDF_TARGET_ESP32
 #define LCD_HOST    HSPI_HOST
 
@@ -82,23 +84,19 @@ void app_main(void)
 
   printf("Configuration complete!!\n");
 
-  for (int i = 0; i < sizeof(buffer); ++ i) {
-    buffer[i] = 0xfa;
-  }
-
-  spi_transaction_t t = {0};
-  t.tx_buffer = buffer;
-  t.length = sizeof(buffer) * 8;
+  ws2812b_t* drv = ws2812b_init(spi);
+  ws2812b_buffer_t* buffer = ws2812b_new_buffer(5);
 
-  spi_transaction_t* rt;
+  ws2812b_buffer_set_rgb(buffer, 0, 255, 0, 0);
+  ws2812b_buffer_set_rgb(buffer, 1, 0, 255, 0);
+  ws2812b_buffer_set_rgb(buffer, 2, 0, 0, 255);
+  ws2812b_buffer_set_rgb(buffer, 3, 255, 255, 0);
+  ws2812b_buffer_set_rgb(buffer, 4, 255, 0, 255);
 
   while (1) {
-    error = spi_device_queue_trans(spi, &t, portMAX_DELAY);
-    assert(error == ESP_OK);
-    printf("SPI Transaction Equeued!!\n");
-
-    error = spi_device_get_trans_result(spi, &rt, portMAX_DELAY);
-    assert(error == ESP_OK);
-    printf("SPI Transaction Sent!!\n");
+    ws2812b_write_sync(drv, buffer);
+    // TODO(rahm) push this into the sync write, or otherwise make this better
+    // to deal with.
+    vTaskDelay(10 / portTICK_PERIOD_MS);
   }
 }