ws2812.c 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  1. #include "module.h"
  2. #include "lauxlib.h"
  3. #include "lmem.h"
  4. #include "platform.h"
  5. #include <stdlib.h>
  6. #include <string.h>
  7. #include "user_interface.h"
  8. #include "driver/uart.h"
  9. #define CANARY_VALUE 0x32383132
  10. typedef struct {
  11. int canary;
  12. int size;
  13. uint8_t colorsPerLed;
  14. uint8_t values[0];
  15. } ws2812_buffer;
  16. // Init UART1 to be able to stream WS2812 data
  17. // We use GPIO2 as output pin
  18. static void ws2812_init() {
  19. // Configure UART1
  20. // Set baudrate of UART1 to 3200000
  21. WRITE_PERI_REG(UART_CLKDIV(1), UART_CLK_FREQ / 3200000);
  22. // Set UART Configuration No parity / 6 DataBits / 1 StopBits / Invert TX
  23. WRITE_PERI_REG(UART_CONF0(1), UART_TXD_INV | (1 << UART_STOP_BIT_NUM_S) | (1 << UART_BIT_NUM_S));
  24. // Pull GPIO2 down
  25. platform_gpio_mode(4, PLATFORM_GPIO_OUTPUT, PLATFORM_GPIO_FLOAT);
  26. platform_gpio_write(4, 0);
  27. // Waits 10us to simulate a reset
  28. os_delay_us(10);
  29. // Redirect UART1 to GPIO2
  30. // Disable GPIO2
  31. GPIO_REG_WRITE(GPIO_ENABLE_W1TC_ADDRESS, BIT2);
  32. // Enable Function 2 for GPIO2 (U1TXD)
  33. PIN_FUNC_SELECT(PERIPHS_IO_MUX_GPIO2_U, FUNC_U1TXD_BK);
  34. }
  35. // Stream data using UART1 routed to GPIO2
  36. // ws2812.init() should be called first
  37. //
  38. // NODE_DEBUG should not be activated because it also uses UART1
  39. static void ICACHE_RAM_ATTR ws2812_write(uint8_t *pixels, uint32_t length) {
  40. // Data are sent LSB first, with a start bit at 0, an end bit at 1 and all inverted
  41. // 0b00110111 => 110111 => [0]111011[1] => 10001000 => 00
  42. // 0b00000111 => 000111 => [0]111000[1] => 10001110 => 01
  43. // 0b00110100 => 110100 => [0]001011[1] => 11101000 => 10
  44. // 0b00000100 => 000100 => [0]001000[1] => 11101110 => 11
  45. // Array declared as static const to avoid runtime generation
  46. // But declared in ".data" section to avoid read penalty from FLASH
  47. static const __attribute__((section(".data._uartData"))) uint8_t _uartData[4] = { 0b00110111, 0b00000111, 0b00110100, 0b00000100 };
  48. uint8_t *end = pixels + length;
  49. do {
  50. uint8_t value = *pixels++;
  51. // Wait enough space in the FIFO buffer
  52. // (Less than 124 bytes in the buffer)
  53. while (((READ_PERI_REG(UART_STATUS(1)) >> UART_TXFIFO_CNT_S) & UART_TXFIFO_CNT) > 124);
  54. // Fill the buffer
  55. WRITE_PERI_REG(UART_FIFO(1), _uartData[(value >> 6) & 3]);
  56. WRITE_PERI_REG(UART_FIFO(1), _uartData[(value >> 4) & 3]);
  57. WRITE_PERI_REG(UART_FIFO(1), _uartData[(value >> 2) & 3]);
  58. WRITE_PERI_REG(UART_FIFO(1), _uartData[(value >> 0) & 3]);
  59. } while(pixels < end);
  60. }
  61. // Lua: ws2812.write("string")
  62. // Byte triples in the string are interpreted as G R B values.
  63. //
  64. // ws2812.init() should be called first
  65. //
  66. // ws2812.write(string.char(0, 255, 0)) sets the first LED red.
  67. // ws2812.write(string.char(0, 0, 255):rep(10)) sets ten LEDs blue.
  68. // ws2812.write(string.char(255, 0, 0, 255, 255, 255)) first LED green, second LED white.
  69. static int ws2812_writegrb(lua_State* L) {
  70. size_t length;
  71. const char *values = luaL_checklstring(L, 1, &length);
  72. // Send the buffer
  73. ws2812_write((uint8_t*) values, length);
  74. return 0;
  75. }
  76. // Handle a buffer where we can store led values
  77. static int ws2812_new_buffer(lua_State *L) {
  78. const int leds = luaL_checkint(L, 1);
  79. const int colorsPerLed = luaL_checkint(L, 2);
  80. luaL_argcheck(L, leds > 0, 1, "should be a positive integer");
  81. luaL_argcheck(L, colorsPerLed > 0, 2, "should be a positive integer");
  82. // Allocate memory
  83. size_t size = sizeof(ws2812_buffer) + colorsPerLed*leds*sizeof(uint8_t);
  84. ws2812_buffer * buffer = (ws2812_buffer*)lua_newuserdata(L, size);
  85. // Associate its metatable
  86. luaL_getmetatable(L, "ws2812.buffer");
  87. lua_setmetatable(L, -2);
  88. // Save led strip size
  89. buffer->size = leds;
  90. buffer->colorsPerLed = colorsPerLed;
  91. // Store canary for future type checks
  92. buffer->canary = CANARY_VALUE;
  93. return 1;
  94. }
  95. static int ws2812_buffer_fill(lua_State* L) {
  96. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  97. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  98. // Grab colors
  99. int i, j;
  100. int * colors = luaM_malloc(L, buffer->colorsPerLed * sizeof(int));
  101. for (i = 0; i < buffer->colorsPerLed; i++)
  102. {
  103. colors[i] = luaL_checkinteger(L, 2+i);
  104. }
  105. // Fill buffer
  106. uint8_t * p = &buffer->values[0];
  107. for(i = 0; i < buffer->size; i++)
  108. {
  109. for (j = 0; j < buffer->colorsPerLed; j++)
  110. {
  111. *p++ = colors[j];
  112. }
  113. }
  114. // Free memory
  115. luaM_free(L, colors);
  116. return 0;
  117. }
  118. static int ws2812_buffer_fade(lua_State* L) {
  119. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  120. const int fade = luaL_checkinteger(L, 2);
  121. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  122. luaL_argcheck(L, fade > 0, 2, "fade value should be a strict positive number");
  123. uint8_t * p = &buffer->values[0];
  124. int i;
  125. for(i = 0; i < buffer->size * buffer->colorsPerLed; i++)
  126. {
  127. *p++ /= fade;
  128. }
  129. return 0;
  130. }
  131. static int ws2812_buffer_get(lua_State* L) {
  132. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  133. const int led = luaL_checkinteger(L, 2) - 1;
  134. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  135. luaL_argcheck(L, led >= 0 && led < buffer->size, 2, "index out of range");
  136. int i;
  137. for (i = 0; i < buffer->colorsPerLed; i++)
  138. {
  139. lua_pushnumber(L, buffer->values[buffer->colorsPerLed*led+i]);
  140. }
  141. return buffer->colorsPerLed;
  142. }
  143. static int ws2812_buffer_set(lua_State* L) {
  144. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  145. const int led = luaL_checkinteger(L, 2) - 1;
  146. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  147. luaL_argcheck(L, led >= 0 && led < buffer->size, 2, "index out of range");
  148. int type = lua_type(L, 3);
  149. if(type == LUA_TTABLE)
  150. {
  151. int i;
  152. for (i = 0; i < buffer->colorsPerLed; i++)
  153. {
  154. // Get value and push it on stack
  155. lua_rawgeti(L, 3, i+1);
  156. // Convert it as int and store them in buffer
  157. buffer->values[buffer->colorsPerLed*led+i] = lua_tonumber(L, -1);
  158. }
  159. // Clean up the stack
  160. lua_pop(L, buffer->colorsPerLed);
  161. }
  162. else if(type == LUA_TSTRING)
  163. {
  164. size_t len;
  165. const char * buf = lua_tolstring(L, 3, &len);
  166. // Overflow check
  167. if( buffer->colorsPerLed*led + len > buffer->colorsPerLed*buffer->size )
  168. {
  169. return luaL_error(L, "string size will exceed strip length");
  170. }
  171. memcpy(&buffer->values[buffer->colorsPerLed*led], buf, len);
  172. }
  173. else
  174. {
  175. int i;
  176. for (i = 0; i < buffer->colorsPerLed; i++)
  177. {
  178. buffer->values[buffer->colorsPerLed*led+i] = luaL_checkinteger(L, 3+i);
  179. }
  180. }
  181. return 0;
  182. }
  183. static int ws2812_buffer_size(lua_State* L) {
  184. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  185. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  186. lua_pushnumber(L, buffer->size);
  187. return 1;
  188. }
  189. static int ws2812_buffer_write(lua_State* L) {
  190. ws2812_buffer * buffer = (ws2812_buffer*)lua_touserdata(L, 1);
  191. luaL_argcheck(L, buffer && buffer->canary == CANARY_VALUE, 1, "ws2812.buffer expected");
  192. // Send the buffer
  193. ws2812_write(buffer->values, buffer->colorsPerLed*buffer->size);
  194. return 0;
  195. }
  196. static const LUA_REG_TYPE ws2812_buffer_map[] =
  197. {
  198. { LSTRKEY( "fade" ), LFUNCVAL( ws2812_buffer_fade )},
  199. { LSTRKEY( "fill" ), LFUNCVAL( ws2812_buffer_fill )},
  200. { LSTRKEY( "get" ), LFUNCVAL( ws2812_buffer_get )},
  201. { LSTRKEY( "set" ), LFUNCVAL( ws2812_buffer_set )},
  202. { LSTRKEY( "size" ), LFUNCVAL( ws2812_buffer_size )},
  203. { LSTRKEY( "write" ), LFUNCVAL( ws2812_buffer_write )},
  204. { LSTRKEY( "__index" ), LROVAL ( ws2812_buffer_map )},
  205. { LNILKEY, LNILVAL}
  206. };
  207. static const LUA_REG_TYPE ws2812_map[] =
  208. {
  209. { LSTRKEY( "write" ), LFUNCVAL( ws2812_writegrb )},
  210. { LSTRKEY( "newBuffer" ), LFUNCVAL( ws2812_new_buffer )},
  211. { LSTRKEY( "init" ), LFUNCVAL( ws2812_init )},
  212. { LNILKEY, LNILVAL}
  213. };
  214. int luaopen_ws2812(lua_State *L) {
  215. // TODO: Make sure that the GPIO system is initialized
  216. luaL_rometatable(L, "ws2812.buffer", (void *)ws2812_buffer_map); // create metatable for ws2812.buffer
  217. return 0;
  218. }
  219. NODEMCU_MODULE(WS2812, "ws2812", ws2812_map, luaopen_ws2812);