; to disassmble:
; avr-objdump -m avr -D firmware.hex
; firmware.hex is located in .pioenvs/uno in project dir

#include "atmega328p_def_inc.S"

.equ LED_DATA_PIN, 7
.equ SRAM_DATA_BUF_ADDR, 255

.equ LED_COUNT, (1024 * 3)
.equ SRAM_DATA_BUF_ADDR_END, SRAM_DATA_BUF_ADDR + (LED_COUNT / 2) + 1

.text

.org 0
  rjmp main

main:
  ; reset system status register
  ldi r16,0
  out SREG,r16

  ; init stack pointer
  ldi r16,lo8(RAMEND)
  out SPL,r16
  ldi r16,hi8(RAMEND)
  out SPH,r16

  ; set data pin as output
  ldi r16, (1 << LED_DATA_PIN)
  out DDRD, r16
  rcall uart_init
  rjmp mainloop


  ;-----------------------------------------------------------
  ; UART
  ;-----------------------------------------------------------
  uart_init:
    ; set baud rate to 1 Mbps
    ldi r16, 1<<U2X0
    sts UCSR0A, r16
    ldi r16, 0
    sts UBRR0H, r16
    ldi r16, 1
    sts UBRR0L, r16

    ; enable receive and transmit, without interrupts
    ldi r16, (1<<RXEN0)|(1<<TXEN0)
    sts UCSR0B,r16

    ; frame format: 8 data bits, 1 stop bit, no parity.
    ldi r16, (1<<UCSZ01)|(1<<UCSZ00)
    sts UCSR0C,r16

    ret

  ; sends one byte over the serial line.
  ; data must reside in r17.
  uart_send_byte:
    ; wait for send register to become empty
    lds r16, UCSR0A
    sbrs r16, UDRE0
    rjmp uart_send_byte

    ; write data to transmit register
    sts UDR0, R17

    ret

  ; receives one byte from the serial line.
  ; data received is moved to r17.
  uart_receive_byte:
    ; wait for receive complete flag
    lds r16, UCSR0A
    sbrs r16, RXC0
    rjmp uart_receive_byte

    ; write data to r17
    lds r17, UDR0

    ret

;-----------------------------------------------------------
; main loop
;-----------------------------------------------------------
mainloop:
  call uart_receive_frame
  call write_leds
  rjmp mainloop


;-----------------------------------------------------------
; receive frame
;-----------------------------------------------------------
; receive a complete frame through serial line
; frame consists of 1 nibble per led, followed by 0xf
; r30 + r31: frame data pointer
; r20 + r21: data counter
; r22: 0x0f
; r23: 0xf0
uart_receive_frame:
  ; reset data pointer
  ldi r30, lo8(SRAM_DATA_BUF_ADDR)
  ldi r31, hi8(SRAM_DATA_BUF_ADDR)

  ; reset data counter
  ldi r20, lo8(LED_COUNT)
  ldi r21, hi8(LED_COUNT)

rcv_loop:
  ; fetch data from uart
  call uart_receive_byte
  st Z+, r17

  ; look for 0xf nibble


  ; decrement counter
  dec r20
  cpi r20, 0
  breq dec_hi
  rjmp rcv_loop

  dec_hi:
  dec r21
  cpi r21, 0
  breq rcv_done
  ldi r20, 0xff
  rjmp rcv_loop

rcv_done:
  ret

; ws2812b led data output
;
; r30 + r31: frame data pointer
; r20 + r21: data counter
; r22: bit counter
; r23: current frame byte
; r26: 3
; r27: 7
write_leds:
  ; reset data pointer
  ldi r30, lo8(SRAM_DATA_BUF_ADDR) ; ZL
  ldi r31, hi8(SRAM_DATA_BUF_ADDR) ; ZH

  ; reset data counter
  ldi r20, lo8(LED_COUNT)
  ldi r21, hi8(LED_COUNT)

  ; reset bit counter
  ldi r22, 0

  ; set constants
  ldi r26, 3
  ldi r27, 7

  ; load first led byte
  ld r23, Z+

  lsl r23
  inc r22

write_nibble_start:

  ; start with two zeros
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  nop
  nop ; +8
  nop
  nop ; +10
  nop
  nop ; +12
  nop
  nop ; +14
  nop ; +15
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  nop
  nop ; +8
  nop
  nop ; +10
  nop
  nop ; +12
  brbs SREG_C, write_one ; +14 if true / +13 if false
  rjmp write_zero ; +15

write_one:
  nop ; +15
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  inc r22
  lsl r23 ; + 8
  cp r26, r22
  breq write_nibble_end3 ; +11 true, +10 false
  cp r27, r22 ; + 11
  breq write_nibble_end4 ; +13 true, +12 false
  nop ; +13
  nop ; +14
  nop ; +15 cycles total high
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  brbs SREG_C, write_one ; +5 true, +4 false
  rjmp write_zero ; + 6

; END WRITE_ONE


write_zero:
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  inc r22
  lsl r23 ; + 8
  cp r26, r22
  breq write_nibble_end1 ; +11 true, +10 false
  cp r27, r22 ; + 11
  breq write_nibble_end2 ; +13 true, +12 false
  nop ; +13
  nop ; +14
  nop ; +15
; END WRITE_ZERO


; end with two zeros
write_nibble_end3:
  nop
  nop
write_nibble_end4:
  nop
  nop
  cbi PORTD, LED_DATA_PIN
write_nibble_end1:
  nop
  nop
write_nibble_end2:
  nop ; +14
  nop ; +15
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  cp r22, r27  ; +5
  breq next_byte ; +7 true, +6 false
  nop ; +7
  nop ; +8
  rjmp next_led ; + 10

next_byte:  ; load next byte
  ld r23, Z+ ; +9
  clr r22 ; + 10

next_led: ; test ledcount max
  cp r31, lo8(LED_COUNT) ; +11
  breq cmp_hi            ; +13 true, +12 false
  rjmp init_next_led     ; +14

cmp_hi:
  nop ; +14
  nop ; +15
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  cp r30, hi8(LED_COUNT) ; +1
  breq done ; +2
  lsl r23 ; +3
  inc r22 ; +4
  nop
  nop ; +6
  nop
  nop ; +8
  nop
  nop ; +10
  nop
  nop ; +12
  nop
  nop ; +14
  nop
  rjmp write_nibble_start

init_next_led:
  nop  ; +15
  sbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  cbi PORTD, LED_DATA_PIN ; +2
  nop
  nop ; +4
  nop
  nop ; +6
  nop
  nop ; +8
  nop
  nop ; +10
  nop
  lsl r23 ; +12
  inc r22
  rjmp write_nibble_start ; +15

done:
  ret