; Listing5-7.asm
;
; A program that demonstrates thread-local storage.

            option     casemap:none
            include    aoalib.inc   ;AoA library + constants
            includelib aoalib.lib   ;Link in aoalib library
            include    c:\masm32\include64\win64.inc
            include    c:\masm32\include64\kernel32.inc
            includelib c:\masm32\lib64\kernel32.lib  
            
            .const

; Program title:

            align   word
ttlStr      byte    "Listing 5-7", 0

            .data
            align   qword 

; The following holds the index returned by TlsAlloc:

tlsIndex    dword   ?

; The following two events synchronize the execution
; of the threads with the main program:

hEvent      qword   2 dup (?)

; Used to protect print statements in this code:

printCS     CRITICAL_SECTION        {}

            .code

; This is the TLS data structure (memory block)
; to allocate on a per-thread basis.

tlsVars_t   struct
v1          dword   ?
v2          dword   ?
v3          dword   ?
v4          dword   ?
tlsVars_t   ends


            option     prologue:none
            
; Utility functions used by the thread procedure.
;
; initBlock-
;
;           ECX contains an event handle on entry.
;           This procedure signals that event on
;           exit. It also uses the binary bits of
;           the handle as an integer to use when
;           initializing the tlsVars_t block.

initBlock   proc
            push    rcx     ;Also keeps stack 16-byte aligned
            
; Retrieve the pointer to the tlsVars_t data for this
; thread:

            mov     ecx, tlsIndex
            call    TlsGetValue
            test    rax, rax        ;Check for error
            jz      badTLS          ;Pointers are never 0
            
; RCX is a 32-bit integer value specifying an initial
; integer to store into the fields of tlsVars_t. Increment
; ECX for each element.

            mov     ecx, [rsp]      ;Retrieve value
            mov     (tlsVars_t ptr [rax]).v1, ecx
            inc     ecx
            mov     (tlsVars_t ptr [rax]).v2, ecx
            inc     ecx
            mov     (tlsVars_t ptr [rax]).v3, ecx
            inc     ecx
            mov     (tlsVars_t ptr [rax]).v4, ecx
            jmp     allDone

            
; If there was an error retrieving the TLS entry, drop
; down here and complain:
                    
badTLS:     call    GetLastError    ;Retrieve the error code
            push    rax
            push    rax             ;Keep stack 16-byte aligned
            lea     rcx, printCS
            mov     edx, INFINITE
            call    EnterCriticalSection
            pop     rdx             ;Retrieve error code
            pop     rdx
            call    print
            byte    "initBlock: bad TLS index: %d", nl, 0
            lea     rcx, printCS
            call    LeaveCriticalSection
            
allDone:    pop     rcx
            ret     
initBlock   endp



; dispBlock-
;
;           Displays the values in the tlsVars_t block
;           associated with this thread:

dispBlock   proc
            push    rcx     ;Also keeps stack 16-byte aligned

; Protect the call to print. Technically, this should be
; called immediately before the call to print; putting it
; here prevents it from messing up register values
; returned by TlsGetValue, etc.

            lea     rcx, printCS
            mov     edx, INFINITE
            call    EnterCriticalSection
            
; Retrieve the pointer to the tlsVars_t data for this
; thread:

            mov     ecx, tlsIndex
            call    TlsGetValue
            test    rax, rax        ;Check for error
            jz      badTLS          ;Pointers are never 0
            
            mov     edx, [rsp]      ;Fetch "thread ID"
            call    print
            byte    "Thread(%d): ", 0
            
            mov     edx,  (tlsVars_t ptr [rax]).v1
            mov     r8d,  (tlsVars_t ptr [rax]).v2
            mov     r9d,  (tlsVars_t ptr [rax]).v3
            mov     r10d, (tlsVars_t ptr [rax]).v4
            call    print
            byte    "dispBlock: %d, %d, %d, %d", nl, 0
            jmp     allDone


; If there was an error retrieving the TLS entry, drop
; down here and complain:
                    
badTLS:     call    GetLastError    ;Retrieve the error code
            push    rax
            push    rax             ;Keep stack 16-byte aligned
            lea     rcx, printCS
            mov     edx, INFINITE
            call    EnterCriticalSection
            pop     rdx             ;Retrieve error code
            pop     rdx
            call    print
            byte    "dispBlock: bad TLS index: %d", nl, 0

allDone:    lea     rcx, printCS
            call    LeaveCriticalSection
            pop     rcx
            ret     
dispBlock   endp

parmOfs     =       16
rbxOfs      =       -8
tlsVarsOfs  =       rbxOfs-sizeof( tlsVars_t )

parm        textequ <qword ptr [rbp+parmOfs]>
saveRBX     textequ <qword ptr [rbp+rbxOfs]>
tlsVars     textequ <tlsVars_t ptr [rbp+tlsVarsOfs]>

threadProc  proc
            push    rbp
            mov     rbp, rsp
            sub     rsp, 128
            mov     saveRBX, rbx
            
; Save the parameter (passed to us in R/ECX) into parm:

            mov     parm, rcx
            
; Display a message indicating we've entered this thread:

            lea     rcx, printCS
            mov     edx, INFINITE
            call    EnterCriticalSection
            mov     edx, dword ptr parm
            call    print
            byte    "Thread(%d): started", nl, 0
            lea     rcx, printCS
            call    LeaveCriticalSection
            
; Upon entry into the thread, call TlsSetValue to store the
; address of the thread-local block into the TLS array:

            mov     ecx, tlsIndex   ;Index was set up by main program
            lea     rdx, tlsVars    ;Allocated block on stack
            call    TlsSetValue

; Here's some code that will initialize the thread local storage
; block with values based on the argument passed to this thread
; procedure.
; Note that this code was moved to a separate procedure to better
; simulate what would happen in real-world code.
            
            mov     rcx, parm
            call    initBlock
            
; Pretend a lot of work is taking place and this thread
; gets pre-empted:

            mov     ecx, 1000
            call    Sleep

; Here's some code that will simply display the values in the block.
; Note that this code was moved to a separate procedure to better
; simulate what would happen in real-world code.
            
            mov     rcx, parm
            call    dispBlock
            
            mov     rbx, saveRBX
            leave
            mov     rcx, parm       ;Get event handle
            call    SetEvent        ;Tell main we're done.
            ret
threadProc  endp
        
; Here is the main assembly language function.

            
            public  asmMain
asmMain     proc
            
            push    rbp
            mov     rbp, rsp
            sub     rsp, 64
            mov     [rbp-8], rbx
            
; Initialize the critical section that will protect
; calls to print in this code:

            lea     rcx, printCS
            call    InitializeCriticalSection
            
; Set up a couple of events so the main program can
; wait for the threads to finish:

            xor     ecx, ecx        ;No security attributes
            xor     edx, edx        ;Auto-reset mode
            xor     r8d, r8d        ;Initial: unsignaled
            xor     r9, r9          ;No event ID; local
            call    CreateEvent
            mov     hEvent[0*8], rax
                        
            xor     ecx, ecx        ;No security attributes
            xor     edx, edx        ;Auto-reset mode
            xor     r8d, r8d        ;Initial: unsignaled
            xor     r9, r9          ;No event ID; local
            call    CreateEvent
            mov     hEvent[1*8], rax
            

; No threads yet, so no need to protect this:
            
            call    print
            byte    "Creating Thread1", nl, 0
            
; Create the first thread:

            xor     rcx, rcx        ;No security attributes
            xor     rdx, rdx        ;Default stack size
            lea     r8, threadProc  ;Address of thread code
            mov     r9, hEvent[0*8] ;Thread parameter
            xor     rax, rax        ;Default thread flags
            mov     [rsp+32], rax   ;Must pass on stack.
            mov     [rsp+40], rax   ;Don't save threadID.
            call    CreateThread
            
; Need to protect call to print here as the thread
; just created could be calling it, too.

            lea     rcx, printCS
            mov     edx, INFINITE
            call    EnterCriticalSection
            call    print
            byte    "Creating Thread2", nl, 0
            lea     rcx, printCS
            call    LeaveCriticalSection
            
; Create the second thread:

            xor     rcx, rcx        ;No security attributes
            xor     rdx, rdx        ;Default stack size
            lea     r8, threadProc  ;Address of thread code
            mov     r9, hEvent[1*8] ;Thread parameter
            xor     rax, rax        ;Default thread flags
            mov     [rsp+32], rax   ;Must pass on stack.
            mov     [rsp+40], rax   ;Don't save threadID.
            call    CreateThread

; Wait for the two threads to complete execution:

            mov     ecx, 2
            lea     rdx, hEvent     ;Array of 2 qwords
            mov     r8d, 1          ;Wait for all signals
            mov     r9d, INFINITE   
            call    WaitForMultipleObjects
            
                        
            
; Terminate the program.

allDone:    
            lea     rcx, printCS
            call    DeleteCriticalSection

            mov     rbx, [rbp-8]
            leave
            ret     ;Returns to caller
        
asmMain     endp
            end