; WTFStopW : a simple stopwatch
; Copyright (C) 2018 Weber Yann
; 
; This program is free software; you can redistribute it and/or modify
; it under the terms of the GNU General Public License as published by
; the Free Software Foundation; either version 3 of the License, or
; any later version.
; 
; This program is distributed in the hope that it will be useful,
; but WITHOUT ANY WARRANTY; without even the implied warranty of
; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
; GNU General Public License for more details.
; 
; You should have received a copy of the GNU General Public License
; along with this program.  If not, see <http://www.gnu.org/licenses/>.

;
; A simple precise stopwatch
; Build : nasm -felf64 wtfstopw.asm && ld -s -melf_x86_64 wtfstopw.o -o wtfstopw
; Build Debug : nasm -g -F dwarf -felf64 -l wtfstopw.lst wtfstopw.asm && ld -melf_x86_64 wtfstopw.o -o wtfstopw
;
; ./wtfstopw [-h] [-r NDIGITS]
; Options :
; 	-h print this help and exit
; 	-r number of digits bellow seconds to display, default is 2
; Interact :
; 	press enter to exit
; 	send SIGINT (with kill -2 ir ctrl + c) for a new lap



[bits 64]

STRUC TIMESPEC_STRUC
	.tv_sec: resq 1
	.tv_nsec: resq 1
ENDSTRUC


%macro TIMESPEC 3
	%1: ISTRUC TIMESPEC_STRUC
		at TIMESPEC_STRUC.tv_sec, dq %2
		at TIMESPEC_STRUC.tv_nsec, dq %3
	IEND
	%define %1.tv_sec %1+TIMESPEC_STRUC.tv_sec
	%define %1.tv_nsec %1+TIMESPEC_STRUC.tv_nsec
%endmacro

STRUC SIGACTION_STRUC
	.sa_handler:      resq      1
	.sa_flags:        resq      1
	.sa_restorer:     resq      1
	.sa_mask:         resb      128
ENDSTRUC

section .data

	sigaction: ISTRUC SIGACTION_STRUC
		at  SIGACTION_STRUC.sa_handler, dq 0
		at  SIGACTION_STRUC.sa_flags, dq 0
		at  SIGACTION_STRUC.sa_restorer, dq 0
		at  SIGACTION_STRUC.sa_mask, times 128 db 0
	IEND
	%define sigaction.sa_handler sigaction+SIGACTION_STRUC.sa_handler
	%define sigaction.sa_flags sigaction+SIGACTION_STRUC.sa_flags
	%define sigaction.sa_restorer sigaction+SIGACTION_STRUC.sa_restorer
	%define sigaction.sa_mask sigaction+SIGACTION_STRUC.sa_mask

	TIMESPEC ts_start, 0, 0
	TIMESPEC ts_cur, 0, 0
	TIMESPEC ts_sleep, 0, 100000000 ; set before mainloop


	faultmsg: db "Fault !", 0xA
	faultmsglen: equ $ - faultmsg

	startmsg: db "Press Enter or ctrl+d to exit and ctrl+c for new lap."
	db 0xA
	startmsglen: equ $ - startmsg

	usage_pre: db "Usage : "
	usage_prelen: equ $ - usage_pre

	usage_post: db " [-h] [-r NDIGITS]", 0xA
		db "Options :", 0xA
		db 0x9, "-h print this help and exit", 0xA
		db 0x9, "-r Number of digits bellow seconds, between [1..8]"
		db 0xA
		db "Interactions :", 0xA
		db 0x9, "Press enter or close stdin to exit", 0xA
		db 0x9, "Send SIGINT (with kill -2 or ctrl + c) for a new lap"
		db 0xA
		db 0x9, "Elapsed time is sent on stderr and laps infos on stdout"
		db 0xA
	usage_postlen: equ $ - usage_post

	badarg_msg: db "Unexpected argument : "
	badarg_msglen: equ $ - badarg_msg

	badval_msg: db "Value for -r should be in [1..8] but got "
	badval_msglen: equ $ - badval_msg

	hours: db "000000000"
	timestr: db ":00:00.0           ", 0x0a
	timestrlen: equ $ - timestr

	time_res: dq 2 ; 2 digits bellow seconds can grow to 8

	nl: db 0x0A
	buf: db 0

	lapsmsg: db 0x0d, "Lap : "
	lapsmsglen: equ $ - lapsmsg

	lapcount: db "00000000"
	lapcountlen: equ $ - lapcount
	laplen: dq 2

	fcntl_flag: dq 0

section .text
global _start
_start:

; parse arguments and set time_res value
jmp arg_parse
arg_ok:

; set stdin non blocking
xor rdx, rdx
xor rdi, rdi
mov rax, 72 ; fcntl
mov rsi, 3 ; F_GETFL
syscall
mov [fcntl_flag], rax
mov rdx, rax
or rdx, 0x800 ; O_NONBLOCK
mov rax, 72 ; fcntl
mov rsi, 4 ; F_SETFL
syscall
cmp rax, 0
jne fault

; initializing lapptr

; preparing SIGINT catch
mov rax, proc_lap_handler
mov qword [sigaction.sa_handler], rax
mov eax, 0x14000000 ; SA_RESTART | SA_RESTORER
mov dword [sigaction.sa_flags], eax
mov rax, sig_restorer
mov qword [sigaction.sa_restorer], rax

mov rax, 13 ; sys_rt_sigaction
mov rdi, 2 ; SIGINT
mov rsi, sigaction
mov rdx, 0 ; NULL
mov r10, 8 ; sig_size
syscall
cmp rax, 0
jne fault

mov rax, 228 ; clock_gettime
mov rdi, 0 ; CLOCK_REALTIME
mov rsi, ts_start
syscall

mov rax, 1 ; write
mov rdi, 2 ; stderr
mov rsi, startmsg
mov rdx, startmsglen
syscall

; set value for ts_sleep.tv_nsec given time_res
; div sleep time by 10 for each digits added bellow seconds
mov rax, 100000000
mov r8, [time_res]
mov r9, 10
xor rdx, rdx
setsleep_loop:
	cmp r8, 1
	jle setsleep_endloop
	div r9
	sub r8, 1
	jmp setsleep_loop
setsleep_endloop:
mov [ts_sleep.tv_nsec], rax

main_loop:
	push 2 ; stderr
	push 0x0D ; \r
	call proc_print_time

	; Attempt to read from stdin
	; if something read, enter has been pressed
	mov rax, 0
	mov rdi, 0
	mov rsi, buf
	mov rdx, 1
	syscall
	cmp rax, 0
	jge flush_stdin ; flush stdin and exit

	mov rax, 35 ; nanosleep
	mov rdi, ts_sleep
	mov rsi, 0
	syscall

jmp main_loop ; main_loop

flush_stdin:
	mov rax, 0
	mov rdi, 0
	mov rsi, buf
	mov rdx, 1
	syscall
	cmp rax, 0
	je newline_exit
	jg flush_stdin

mov rdi, 0 ; EXIT OK
;
; Expect rdi to be the return code
;
exit:
	push rdi
	; restoring stdin state
	mov rax, 72 ; fcntl
	xor rdi, rdi
	mov rsi, 4 ; F_SETFL
	mov rdx, [fcntl_flag]
	syscall
	cmp rax, 0
	je exit_end
	pop rdi ; failed to restore
	push 1 ; exit FAIL

	exit_end:
	mov rax, 60 ; sys_exit
	pop rdi ; return code
	syscall
fault:
	mov rax, 1 ; write
	mov rdi, 2 ; stderr
	mov rsi, nl
	mov rdx, 1
	syscall
	mov rax, 1
	mov rsi, faultmsg
	mov rdx, faultmsglen
	syscall

	mov rdi, 1 ; failure
	jmp exit

newline_exit:
	mov rax, 1
	mov rdi, 1
	mov rsi, nl
	mov rdx, 1
	syscall

	mov rdi, 0 ; exit OK
	jmp exit

;
; Print current time on FD and add a leading char CHR
; push FD & push CHR to set arguments
;
proc_print_time:
	; push ret addr before arguments
	pop r8
	pop r9
	pop r10
	push r8
	push r10
	push r9

	; updating ts_cur time
	mov rax, 228 ; clock_gettime
	mov rdi, 0 ; CLOCK_REALTIME
	mov rsi, ts_cur
	syscall

	; Calculating elapsed ns
	mov rax, [ts_cur.tv_nsec]
	mov rbx, [ts_start.tv_nsec]
	sub rax, rbx
	cmp rax, 0
	jge print_time_us_cont
	; negativ result
	add rax, 1000000000
	mov rbx, [ts_cur.tv_sec]
	sub rbx, 1
	mov [ts_cur.tv_sec], rbx
	print_time_us_cont:
	
	; Divide result given time_res
	mov r10, rax
	mov rax, 1000000000
	mov r9, [time_res]
	mov r8, 10
	xor rdx, rdx
	print_time_respow:
		div r8
		sub r9, 1
		cmp r9, 0
		jg print_time_respow
	mov rcx, rax
	mov rax, r10
	xor rdx, rdx
	div rcx

	; set the us char in timestr
	mov r8, timestr
	mov r9, [time_res] ; r9 count the number of digits
	add r8, 6 ; first digits after seconds
	add r8, r9 ; r8 points on last char before \r
	print_time_us_loop:
		xor rdx, rdx
		mov rcx, 10
		div rcx
		add dl, 0x30
		mov [r8], dl
		sub r8, 1
		sub r9, 1
		cmp r9, 0
		jg print_time_us_loop

	; handling seconds, minutes & hours
	mov rax, [ts_cur.tv_sec]
	mov rbx, [ts_start.tv_sec]
	sub rax, rbx
	; rax now contain elapsed seconds
	; filling timestr with seconds & minutes

	xor rdx, rdx
	mov rcx, 10
	div rcx
	add dl, 0x30
	mov byte [timestr + 5], dl

	xor rdx, rdx
	mov rcx, 6
	div rcx
	push rax
	add dl, 0x30
	mov byte [timestr + 4], dl

	xor rdx, rdx
	mov rcx, 10
	div rcx
	add dl, 0x30
	mov byte [timestr + 2], dl

	pop rax
	xor rdx, rdx
	mov rcx, 6
	div rcx
	add dl, 0x30
	mov byte[timestr + 1], dl

	; filling the hours buffer
	; r8 will contain our digits counter : max is 8
	mov r8, 8
	print_time_hours_loop:
		mov rcx, 10
		xor rdx, rdx
		div rcx
		add dl, 0x30
		cmp rax, 0
		jne print_time_hours_print_mod
		cmp rdx, 0
		je print_time_hours_loop_end

		print_time_hours_print_mod:
		mov r9, hours
		add r9, r8
		mov byte [r9], dl
		cmp r8, 0
		je fault
		sub r8, 1
		cmp rax, 0
		jne print_time_hours_loop
	print_time_hours_loop_end:

	; print hours + timestr
	add r8, 1
	cmp r8, 7
	jle print_time_hours_cont
	mov r8, 7 ; maximum value for r8
	print_time_hours_cont:
	mov r9, hours
	add r9, r8
	mov rcx, 9
	sub rcx, r8 ; rcx is hours size 
	add rcx, timestrlen ; add to timestrlen

	; Set leading char
	pop r10
	mov byte [timestr+timestrlen-1], r10b

	mov rax, 1 ; write
	pop r10 ; print_time FD arg
	mov rdi, r10 ; print_time arg
	mov rsi, r9 ; start hours pointer
	mov rdx, rcx ; size to timestr end
	syscall

	ret

;
;	sig handler for SIGINT displaying lap count and time on stdout
;
proc_lap_handler:
	mov rax, 1
	mov rdi, 1
	mov rsi, lapsmsg
	mov rdx, 5 ; "Lap "
	syscall

	; increment the lapcount str directly
	mov rbx, lapcount ; first digit ptr
	add rbx, lapcountlen ; rightmost digit ptr
	sub rbx, 1
	mov r8, 1 ; counter
	lap_handler_inc_lap:
		mov r10b, [rbx]
		cmp r10b, 0x39 ; '9'
		jl lap_handler_inc_end

		add r8, 1
		cmp r8, [laplen]
		jl lap_handler_laplen_noupd
		mov [laplen], r8 ; update laplen

		lap_handler_laplen_noupd:
		mov byte [rbx], 0x30 ; set current digit to '0'
		sub rbx, 1
		cmp rbx, lapcount
		jl fault
		jmp lap_handler_inc_lap
	lap_handler_inc_end:
		add r10b, 1
		mov [rbx], r10b
	

	mov rax, 1
	mov rdi, 1
	mov rdx, [laplen]
	mov rsi, lapcount
	add rsi, lapcountlen
	sub rsi, rdx ; leftmost digit ptr
	syscall

	mov rax, 1
	mov rdi, 1
	mov rsi, lapsmsg + 4
	mov rdx, 3 ; " : "
	syscall

	push 1 ; stdout
	push 0x0A ; \n
	call proc_print_time

	ret

sig_restorer:
	mov rax, 15 ; sys_rt_sigreturn
	xor rdi, rdi
	syscall

;
;	Argument parsing
;
arg_parse:
	; Checking argument count
	mov rax, [rsp]
	cmp rax, 1
	je arg_ok
	cmp rax, 3
	jle parse_r
	arg_err: ; badval & badarg jmp here too
		mov rax, [rsp+8] ; argv[0] program name
		mov rdi, -1 ; return status
		jmp usage
	parse_r:
	mov rax, [rsp+16] ; 1st arg should be "-r"
	mov bl, [rax]
	cmp bl, '-'
	jne badarg
	mov bl, [rax+1]
	cmp bl, 'h' ; -h
	je arg_err
	cmp bl, 'r'
	jne badarg
	mov bl, [rax+2]
	cmp bl, 0
	jne arg_nxt ; the value seems to be just after the -r like -r2
	nxt_arg:
		; the 1st argument is -r the second must be the time_res
		; check that the arg exists
		mov rax, [rsp]
		cmp rax, 3
		jne arg_err
		mov rax, [rsp+24]
		jmp arg_cont
	arg_nxt:
		; check that there is no more args
		mov rbx, [rsp]
		cmp rbx, 2
		jne arg_err
		add rax, 2
	arg_cont:
	; rax should point on the value
	mov bl, [rax+1]
	cmp bl, 0
	jne badval
	xor rbx, rbx
	mov bl, [rax]
	cmp bl, '1'
	jl badval
	cmp bl, '8'
	jg badval
	sub bl, '0'
	mov [time_res], rbx
	jmp arg_ok

; print an error message, usage and exit with rax pointing the value
badval:
	mov rsi, badval_msg
	mov rdx, badval_msglen
	jmp arg_strerr

; print an error message, usage and exit
badarg:
	mov rsi, badarg_msg
	mov rdx, badarg_msglen
	mov rax, [rsp+16]
	jmp arg_strerr

; rsi msg ptr, rdx msg len, rax arg ptr
arg_strerr:
	mov r8, rax
	mov rax, 1
	mov rdi, 1
	syscall
	mov rsi, r8
	mov rax, r8
	call proc_strlen
	mov rax, 1
	syscall
	mov rax, 1
	mov rsi, nl
	mov rdx, 1
	syscall
	syscall
	jmp arg_err

;
; Print usage and exit
; 	Except rax to point on programm name and rdi to be the return code
;
usage:
	push rdi
	push rax

	mov rax, 1 ; write
	mov rdi, 1 ; stdout
	mov rsi, usage_pre
	mov rdx, usage_prelen
	syscall

	pop rsi
	mov rax, rsi
	call proc_strlen
	mov rax, 1
	syscall

	mov rax, 1
	mov rsi, usage_post
	mov rdx, usage_postlen
	syscall

	mov rax, 60
	pop rdi
	syscall

; With rax pointing on the string, the result will be in rdx
proc_strlen:
	mov r8, rax
	mov r9, r8
	sub r9, 1
	strlen_loop:
		add r9, 1
		mov al, [r9]
		cmp al, 0
		jne strlen_loop
	sub r9, r8
	mov rdx, r9
	ret

.end: