# type MallocBlock = {
#    block-size : U64
#    block-prev : *MallocBlock
#    block-next : *MallocBlock
#    [data] : U8[block-size]
# }

.block-size := λ(: b MallocBlock[]). (: (
   (let r ([]( (as b U64[]) 0_u64 )))
   r
) U64);
.block-prev := λ(: b MallocBlock[]). (: (
   (let r ([]( (as b U64[]) 1_u64 )))
   (as r MallocBlock[])
) MallocBlock[]);
.block-next := λ(: b MallocBlock[]). (: (
   (let r ([]( (as b U64[]) 2_u64 )))
   (as r MallocBlock[])
) MallocBlock[]);
.block-data := λ(: b MallocBlock[]). (: (
   (as (+( (as b U64) 24_u64 )) ?[])
) ?[]);

set.block-size := λ(: b MallocBlock[])(: sz U64). (: (
   (set[]( (as b U64[]) 0_u64 sz ))
) Nil);
set.block-prev := λ(: b MallocBlock[])(: n MallocBlock[]). (: (
   (set[]( (as b U64[]) 1_u64 (as n U64) ))
) Nil);
set.block-next := λ(: b MallocBlock[])(: n MallocBlock[]). (: (
   (set[]( (as b U64[]) 2_u64 (as n U64) ))
) Nil);

malloc-block-head := 0_u64;
malloc-block-tail := 0_u64;
malloc-first-free := 0_u64;
malloc-page-tail := 0_u64;

malloc-new-page := λ. (: (
   (mov( 12_u64 RAX ))
   (mov( 0_u64 RDI ))
   (syscall())

   (mov( RAX R8 )) #current page break in R8
   (set malloc-block-head (as R8 U64))
   (set malloc-block-tail (as R8 U64))

   # 8 GB
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (add( 1073741824_u64 R8 ))
   (mov( 12_u64 RAX ))
   (mov( R8 RDI ))
   (syscall())

   (set malloc-page-tail (as R8 U64))

   # sys_brk initially zeroes memory for each process
   # so the first MallocBlock::header will be zero initialized
) Nil);

realloc := λ(: ptr ?[])(: sz U64). (: (
   (let block (as (-( (as ptr U64) 24_u64 )) MallocBlock[]))
   (let delta (-( (as sz I64) (as (.block-size block) I64) )))

   (if (||(
      (>=( (.block-size block) sz ))
      (<=( delta (as (.free-space block) I64) ))
   )) (
      (set.block-size( block sz ))
      #(if (&&( (==( (as (.block-next block) U64) 0_u64 ))
      #         (<( (.block-size block) sz ))
      #    )) (
      (if (>( delta 0_i64 )) (
         (set malloc-block-tail (+( malloc-block-tail (as delta U64) )))
      ) ())
   ) (
      (let new-data (malloc sz))
      (let new-block (as (-( (as new-data U64) 24_u64 )) MallocBlock[]))
      (let bi 0_u64)
      (while (<( bi sz )) (
         (let dst (as (.block-data new-block) U8[]))
         (let src (as (.block-data block) U8[]))
         (set[]( dst bi ([]( src bi )) ))
         (set bi (+( bi 1_u64 )))
      ))
      (free(.block-data block))
      (set block new-block)
   ))

   (.block-data block)
) ?[]);

.free-space := λ(: head MallocBlock[]). (: (
   (let space 0_u64)
   (if (==( (as (.block-next head) U64) 0_u64 )) (
      (set space 1073741824_u64)
   ) (
      (set space (-(
         (as (.block-next head) U64)
         (+(
            (as head U64)
            (+( 24_u64 (.block-size head) ))
         ))
      )))
   ))
   space
) U64);

free := λ(: ptr ?[]). (: (
   (let block (as (-( (as ptr U64) 24_u64 )) MallocBlock[]))
   (let prev (.block-prev block))
   (let next (.block-next block))
   (if (!=( (as prev U64) 0_u64 )) (
      (set.block-next( prev next ))
   ) ())
   (if (!=( (as next U64) 0_u64 )) (
      (set.block-prev( next prev ))
   ) ())
) Nil);

malloc-large-size := 1024_u64; # need to tune this to find empirical optima

malloc := λ (: sz U64). (: (
   (malloc( sz sz ))
) ?[]);

malloc := λ (: sz U64)(: min-capacity U64). (: (
   (if (==( malloc-block-head 0_u64 )) (malloc-new-page()) ())

   (let scan-head (as malloc-first-free MallocBlock[]))
   (if (==( (as scan-head U64) 0_u64 )) (set scan-head (as malloc-block-head MallocBlock[])) ())

   (let scan-count 0_u64)
   (while (>( (+( min-capacity 24_u64 )) (.free-space scan-head) )) (
      (set scan-head (.block-next scan-head))
      (set malloc-first-free (as scan-head U64))
      (if (>=( min-capacity malloc-large-size )) () (
         (set scan-count (+( scan-count 1_u64 )))
      ))
   ))
   
   (let misalignment (%( malloc-block-tail 8_u64 )))
   (let alignment (%( (-( 8_u64 misalignment )) 8_u64 )))
   (set malloc-block-tail (+( malloc-block-tail alignment )))
   (let block (as malloc-block-tail MallocBlock[]))
   (let data (+( (as block U64) 24_u64 )))
   (set malloc-block-tail (+( data sz )))

   # initialize block
   (let block-next (.block-next scan-head))
   (if (!=( (as block-next U64) 0_u64 )) (
      (set.block-prev( block-next block ))
   ) ())
   (set.block-next( scan-head block ))
   (set.block-size( block sz ))
   (set.block-next( block (as 0_u64 MallocBlock[]) ))
   (if (==( (as block U64) (as scan-head U64) )) (
      (set.block-prev( block (as 0_u64 MallocBlock[]) ))
   ) (
      (set.block-prev( block scan-head ))
   ))

   (as data ?[])
) ?[]);