;=----------------------------------------------------------------------------=;
;                            GNU GPL OS/K                                      ;
;                                                                              ;
;   Desc:   Memory management from protected mode                              ;
;                   (x86_64 architecture only)                                 ;
;                                                                              ;
;                                                                              ;
;   Copyright © 2018-2019 The OS/K Team                                        ;
;                                                                              ;
;   This file is part of OS/K.                                                 ;
;                                                                              ;
;   OS/K 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          ;
;   (at your option) any later version.                                        ;
;                                                                              ;
;   OS/K 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 OS/K.  If not, see <https://www.gnu.org/licenses/>.             ;
;=----------------------------------------------------------------------------=;

%define MAX_MEMORY 4 ; GiB

[BITS 32]
[section .text]
; ---------------------------------------------------------------------------- ;
;    Constructor for the page tables in protected mode                         ;
; ---------------------------------------------------------------------------- ;
Setup_paging:
    ;; Map the first PML4 entry to PDP table
    mov eax, PDP_table
    or eax, 1 << 1 | 1 << 0             ; present + writable
    mov [PML4_table], eax

    ;; Map the PDP entries to PD tables
    mov ebx, PD_table                   ; start address
    mov ecx, 0x0                        ; counter variable
.map_pdp_table:
    mov eax, ebx
    or eax, 1 << 1 | 1 << 0             ; present + writable
    mov [PDP_table + 8 * ecx], eax
    inc ecx
    add ebx, 4096
    cmp ecx, MAX_MEMORY                 ; PDP table is mapped if MAX_MEMORY
    jne .map_pdp_table                  ; else map the next entry

    ;; Map each PD entry to a 'huge' 2MiB page

    mov ecx, 0x0                        ; counter variable
.map_pd_table:
    ;; map ecx-th PD entry to a huge page that starts at address 2MiB*ecx
    mov eax, 0x200000
    mul ecx                             ; start address of ecx-th page
    or eax, 1 << 7 | 1 << 1 | 1 << 0    ; present + writable + huge
    mov [PD_table + ecx * 8], eax
    inc ecx
    cmp ecx, 512 * MAX_MEMORY           ; PD table is mapped if 512
    jne .map_pd_table                   ; else map the next entry

; ---------------------------------------------------------------------------- ;
;    Enable long mode and paging                                               ;
; ---------------------------------------------------------------------------- ;
Go64:
    pusha

    ;; Registering paging
    mov eax, PML4_table
    mov cr3, eax                        ; Load PML4 to cr3

    mov eax, cr4
    or eax, 1 << 5
    mov cr4, eax                        ; Enable PAE

    ;; Activate long mode
    mov ecx, 0xC0000080                 ; Address of MSR
    rdmsr                               ; Read MSR
    or eax, 1 << 8                      ; LME = 1. (Long Mode Enable)
    wrmsr                               ; Write MSR

    ;; Enable paging
    mov eax, cr0
    or eax, 1 << 31                     ; Make MSR bit 31 (PG = Paging) to 1 :
                                        ;    |1|000000000000000000000000000000
                                        ;     |
                                        ;     `------ Paging bit
    mov cr0, eax
    jmp .end
    nop
    nop
.end:

    popa
    ret

; ---------------------------------------------------------------------------- ;
;    Checks if Grub has enabled A20 line                                       ;
; ---------------------------------------------------------------------------- ;
CheckA20:
    pushad
    mov edi, 0x112345                   ; Odd megabyte address.
    mov esi, 0x012345                   ; Even megabyte address.
    mov [esi],esi                       ; Making sure that both addresses contain diffrent values.
    mov [edi],edi                       ;    (if A20 line is cleared the two pointers would point to the address 0x012345 that would contain 0x112345 (edi))
    cmpsd                               ; Compare addresses to see if the're equivalent.
    popad
    jne .A20_on                         ; If not equivalent , A20 line is set.
    jmp .A20_err
.A20_on:
    ret
.A20_err:
    mov ax, "04"                        ; ERROR 04 : A20 line failed
    jmp Error

[BITS 64]
; ---------------------------------------------------------------------------- ;
;    Initilizes the stack                                                      ;
; ---------------------------------------------------------------------------- ;
InitStack:
    xor rax, rax                        ; "Fill pattern"
    push rcx
    push rdi

    ;; Begin address to fill and length
    mov qword [realKernelEnd], kernelEnd
    mov rdi, kernelEnd
    mov rcx, KERNEL_STACK - kernelEnd - 16
                                        ; stop before the return address

    ;; If bit 0 is on, fill one byte
    sar rcx, 1                          ; Shift bit 0 into CY
    jnc $ + 3
    stosb

    ;; We are word aligned and if bit 1 was on fill another word
    sar rcx, 1                          ; Shift bit 1 into CY
    jnc $ + 4
    stosw

    ;; We are dword aligned and if bit 2 was on fill another dword
    sar    rcx, 1                       ; Shift bit 2 into CY
    jnc    $ + 3
    stosd
    ;; RCX now equals the number of qwords to fill
    repnz stosq                      ; Finish by writing RCX qwords.

    pop rdi
    pop rcx
    ret