; 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]
%use smartalign
ALIGNMODE k8

%define O_NONBLOCK 0x800

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


%macro TIMESPEC 1
	%1: ISTRUC TIMESPEC_STRUC
		at TIMESPEC_STRUC.tv_sec, dq 0
		at TIMESPEC_STRUC.tv_nsec, dq 0
	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

; from https://www.linuxnasm.be/home/library/lib-includes/signals-inc
%macro SIGACTION 1
	%1: 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
%endmacro

section .data
align 8

	TIMESPEC ts_start
	TIMESPEC ts_cur
	TIMESPEC ts_sleep

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

	hours: times 8 db '0' ; 8 digits is the max for hours
	timestr: db ":00:00." 
		times 8 db '0' ; 8 digits for nanoseconds
		times 0xE db ' '
		db 0
	timestrend: align 8
	timestrlen: equ timestrend - timestr
	%define HOURSLEN 8

	buf: db 0

	fcntl_flag: dq 0
	SIGACTION sigaction

	lapsmsg: db "Lap : "
	lapsmsglen: equ $ - lapsmsg

	lapcount: times 20 db '0' ; allows storing decimal repr of 1<<64
	lapcountlen: equ $ - lapcount
	laplen: dq 2

	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

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

	nl: db 0x0A
	cr: db 0xd


section .text
global _start
_start:

; parse arguments and set time_res value
call arg_parse

mov rbx, ts_start
%define RBXPTR(x) rbx + x - ts_start

; set stdin reads non blocking
xor rdx, rdx
xor rdi, rdi
mov rax, 72 ; fcntl
mov rsi, 3 ; F_GETFL
syscall
mov [RBXPTR(fcntl_flag)], rax
mov rdx, rax
or dx, O_NONBLOCK
mov rax, 72 ; fcntl
mov rsi, 4 ; F_SETFL
syscall
test al, al
js fault

; preparing SIGINT catch
mov rax, proc_lap_handler
mov qword [RBXPTR(sigaction.sa_handler)], rax
mov ecx, 0x14000000 ; SA_RESTART | SA_RESTORER
mov dword [RBXPTR(sigaction.sa_flags)], ecx
add rax, sig_restorer - proc_lap_handler
mov qword [RBXPTR(sigaction.sa_restorer)], rax

mov rax, 13 ; sys_rt_sigaction
mov rdi, 2 ; SIGINT
mov rsi, rbx
add rsi, sigaction - ts_start
;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, rbx
syscall

mov rax, 1 ; write
mov rdi, 2 ; stderr
mov rsi, rbx
add rsi, startmsg - ts_start
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
xor rcx, rcx
mov cl, [RBXPTR(time_res)]
dec cl
test cl, cl
jz setsleep_loopend

mov rsi, 10
xor rdx, rdx
setsleep_loop:
	div rsi
	loop setsleep_loop
setsleep_loopend:

mov [RBXPTR(ts_sleep.tv_nsec)], rax
mov r12w, 0x0d02
mov r13, proc_print_time
mov r14, buf ; nanosleep
mov r15, ts_sleep

std ; set DF for string operations

align 16
main_loop:

	call r13 ; proc_print_time

	; Attempt to read from stdin (syscall prepared by proc_print_time)
	; if something read, enter has been pressed
	syscall
	test al, al
	jns flush_stdin ; flush & exit
	
	mov rax, 35 ; nanosleep
	mov rdi, r15
	xor rsi, rsi
	syscall

jmp main_loop ; main_loop

flush_stdin:
	jz newline_exit
	xor rax, rax
	;xor rdi, rdi ; done in main_loop
	;mov rsi, r14
	;mov rdx, 1
	syscall
	test al, al
	jns flush_stdin

	xor rbx, rbx ; EXIT OK
	jmp exit

	newline_exit:
	mov rdi, 2 ; stderr
	call proc_nl
	xor rbx, rbx ; exit OK
	jmp exit

;
; Expect rbx to be the return code
;
exit:
	; restoring stdin state
	mov rax, 72 ; fcntl
	xor rdi, rdi
	mov rsi, 4 ; F_SETFL
	mov rdx, [fcntl_flag]
	syscall
	mov rdi, rbx ; return code
	test rax, rax
	cmovnz rdi, rax ; exit FAIL
	mov rax, 60 ; sys_exit
	syscall
fault:
	mov rdi, 2 ; stderr
	call proc_nl
	mov rax, 1
	mov rsi, faultmsg
	mov rdx, faultmsglen
	syscall

	mov rbx, 1 ; failure
	jmp exit
;
; Print current time on FD and add a leading char CHR
; argument are in r12w r12b is FD and r12w(h) is CHR
; 	Warning : DF must be set
;
align 16
proc_print_time:
	
	mov rax, 228 ; clock_gettime
	xor rdi, rdi ; CLOCK_REALTIME
	mov rsi, ts_cur
	syscall ; updating ts_cur time


	mov rsi, ts_start.tv_sec
%define RSIPTR(x) rsi + x - ts_start.tv_sec

	mov rax, [RSIPTR(ts_cur.tv_nsec)]
	sub rax, [ts_start.tv_nsec] ; allign nano sec loop
	jge print_time_us_cont
	add rax, 1000000000 ; negativ result
	sub qword [RSIPTR(ts_cur.tv_sec)], 1
	print_time_us_cont:

	xor rdx, rdx
	div qword [RSIPTR(ts_sleep.tv_nsec)] ; Divide result given time_res

	mov bx, r12w
	shr bx, 8
	mov byte [RSIPTR(timestrend - 1)], bl ; set last chr from 1st argument

	; set the nanosec chars (time_res chars) in timestr
	mov rbx, 0x2020202020202020
	mov rcx, [RSIPTR(time_res)]
	mov rdi, 10
	xor rdx, rdx
	align 16 ; should be allready alligned
	procpt_loopns:
		div rdi
		or dl, 0x30 ; '0'
		shl rbx, 8
		xchg bl, dl
		loop procpt_loopns
	mov [timestr+7], rbx

	; filling timestr with seconds & minutes chars
	mov rax, [RSIPTR(ts_cur.tv_sec)]
	sub rax, [RSIPTR(ts_start.tv_sec)] ; rax now contain elapsed seconds

	mov rcx, 10
	mov rdi, 6

	xor rdx, rdx
	div rcx
	mov bh, dl
	xor dl, dl
	div rdi
	mov bl, dl
	or bx, 0x3030
	mov word [timestr + 4], bx

	xor dl, dl
	div rcx
	mov bh, dl
	xor dl, dl
	div rdi
	mov bl, dl
	or bx, 0x3030
	mov word [timestr + 1], bx

	; using rbx to stores the hours string
	xor rbx, rbx
	mov rcx, HOURSLEN
	mov rdi, 10
	xor dl, dl
	align 16
	procpt_looph:
		div rdi
		shl rbx, 8
		xchg bl, dl
		test rax, rax
		loopnz procpt_looph

	xor rax, rax
	mov al, HOURSLEN - 2
	cmp cx, HOURSLEN - 2
	cmovl ax, cx ; rax stores hours digit count (at least 2)
	test cl, cl
	jz procpt_looph2_end
	align 16
	procpt_looph2: ; end rbx left shift before sto
		shl rbx, 8
		loopne procpt_looph2
	procpt_looph2_end:

	mov rdi, 0x3030303030303030 ; ASCII convertion
	or rbx, rdi
	mov [RSIPTR(hours)], rbx

	xor rdi, rdi
	mov dil, r12b ; fd as 2nd argument
	mov rdx, timestrlen + HOURSLEN
	sub rdx, rax ; timestr + hours len
	lea rsi, [RSIPTR(hours + rax)] ; hours start pointer
	mov rax, 1 ; write
	syscall
	test rax, rax
	js fault

; ALLIGN TRICK : prepare read syscall for main_loop

	xor rax, rax
	xor rdi, rdi
	mov rsi, r14
	mov rdx, 1

	ret
%undef RSIPTR



;
;	sig handler for SIGINT displaying lap count and time on stdout
;
proc_lap_handler:
	mov rax, 1
	mov rdi, 2
	mov rsi, cr
	mov rdx, 1
	syscall ; \r on stderr
	test rax, rax
	js fault

	mov rax, 1
	mov rdi, 1
	mov rsi, lapsmsg
	mov rdx, 4 ; "Lap "
	syscall
	test rax, rax
	js fault

	; increment the lapcount str directly
	mov rcx, lapcountlen - 1; digit counter starting by the right most
	proclap_loop:
		cmp byte [lapcount + rcx], '9'
		jl proclap_loopend
		mov byte [lapcount + rcx], '0' ; set current digit to '0'
		loop proclap_loop
	proclap_loopend:
	add byte [lapcount + rcx], 1 ; increment current digit

	mov rdx, lapcountlen
	sub rdx, rcx
	cmp rdx, [laplen] ; update laplen if needed
	cmovl rdx, [laplen]
	mov [laplen], rdx
	mov rsi, lapcount + lapcountlen
	sub rsi, rdx
	
	mov rax, 1
	mov rdi, 1 ; stdout
	syscall
	test rax, rax
	js fault

	mov rax, 1 ; write
	mov rsi, lapsmsg + 3
	mov rdx, 3 ; " : "
	syscall
	test rax, rax
	js fault

	std ; set DF for string operations
	mov r12w, 0x0A01 ; \n & stdout
	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 rcx, [rsp+8] ; argc
	mov rsi, [rsp+24] ; argv[1]
	mov rdi, badarg_msg

	cmp rcx, 1
	je arg_ok
	cmp rcx, 3
	jle parse_r
	arg_err: ; badval & badarg jmp here too
		mov rax, [rsp+16] ; argv[0] program name
		mov rdi, -1 ; return status
		jmp usage
	parse_r:
	mov rax, rsi
	cmp byte [rsi], '-'
	jne badarg
	mov bl, [rsi+1]
	cmp bl, 'h' ; -h
	je arg_err
	cmp bl, 'r'
	jne badarg
	cmp byte [rsi+2], 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
		cmp cl, 3
		jne arg_err
		mov rsi, [rsp+32]
		jmp arg_cont
	arg_nxt:
		; check that there is no more args
		cmp cl, 2
		jne arg_err
		add rsi, 2
	arg_cont:
	; rsi should point on the ASCII digit counter
	cmp byte [rsi+1], 0
	jne badval
	xor rbx, rbx
	mov bl, [rsi]
	cmp bl, '1'
	jl badval
	cmp bl, '8'
	jg badval
	xor bl, 0x30 ; '0'
	mov [rdi + time_res - badarg_msg], rbx

	arg_ok:
	ret

; print an error message, usage and exit with rax pointing the value and
; rdi badarg_msg ptr
badval:
	mov rbx, rsi

	lea rsi, [rdi + badval_msg - badarg_msg]
	mov rdx, badval_msglen
	jmp arg_strerr

; print an error message, usage and exit
badarg:
	mov rbx, rsi
	mov rsi, rdi
	mov rdx, badarg_msglen
	jmp arg_strerr

; rsi msg ptr, rdx msg len, rbx arg ptr
arg_strerr:
	mov rax, 1
	mov rdi, 1
	syscall
	mov rsi, rbx
	mov rax, rbx
	call proc_strlen
	mov rax, 1
	syscall
	call proc_nl
	call proc_nl
	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
	dec r9
	strlen_loop:
		inc r9
		mov al, [r9]
		cmp al, 0
		jne strlen_loop
	sub r9, r8
	mov rdx, r9
	ret

; with rdi the fd
proc_nl:
	mov rax, 1
	mov rsi, nl
	mov rdx, 1
	syscall
	ret

.end: