<commit_before>open Lwt
open Printf
open V1_LWT
open Core_kernel.Std
open PosMapIntf
open Bin_prot.Std

module Make (MakePositionMap : PosMapF) (BlockDevice : BLOCK) = struct

  module PositionMap = MakePositionMap(BlockDevice)

  type 'a io = 'a Lwt.t

  type page_aligned_buffer = Cstruct.t

  type error = [
    | `Unknown of string
    | `Unimplemented
    | `Is_read_only
    | `Disconnected
    ]

  type superblock = {
      offset : int64;
      length : int;
    } [@@ deriving bin_io]

  type info = {
      read_write: bool;
      sector_size: int;
      size_sectors: int64;
    } [@@ deriving bin_io]

  type structuralInfo = {
      height : int;
      numLeaves : int64;
      sectorsPerBlock : int;
    } [@@ deriving bin_io]

  type core = {
      info : info;
      structuralInfo : structuralInfo;
      bucketSize : int64;
      offset : int64;
      desiredBlockSize : int;
      stash : Stash.t;
    } [@@ deriving bin_io]

  type t = {
      info : info;
      structuralInfo : structuralInfo;
      bucketSize : int64;
      offset : int64;
      desiredBlockSize : int;
      stash : Stash.t;
      positionMap : PositionMap.t;
      blockDevice : BlockDevice.t;
    }

  let core { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash } =
    { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash }

  let extendCore ({ info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash } : core) positionMap blockDevice =
    { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash ; positionMap ; blockDevice}

  let bin_size_t t =
    let coreSize = core t |> bin_size_core in
    let posMapSize = PositionMap.bin_size_t t.positionMap in
    coreSize + posMapSize

  let bin_write_t buf ~pos t =
    let pos = core t |> bin_write_core buf ~pos in
    PositionMap.bin_write_t buf ~pos t.positionMap

  let bin_writer_t = { Bin_prot.Type_class.size = bin_size_t; write = bin_write_t  }

  let bin_read_t buf ~pos_ref blockDevice =
    let core = bin_read_core buf ~pos_ref in
    let positionMap = PositionMap.bin_read_t buf ~pos_ref blockDevice in
    extendCore core positionMap blockDevice

  type id = string

  type op =
    | Read
    | Write

  type bucket = OBlock.t list

  let sizeOfAddressInBytes = 8

  let ( >>= ) x f = x >>= function
                    | `Error e -> return (`Error e)
                    | `Ok x -> f x

  let get_info { info } = return info

  let getStructuralInfo { structuralInfo } = return structuralInfo

  let floor_log x =
    let rec loop acc = function
      | 1L -> acc
      | x -> loop (acc + 1) Int64.(x / 2L)
    in loop 0 x

  let diskSize t =
    Int64.((t.info.size_sectors * of_int t.structuralInfo.sectorsPerBlock) + (PositionMap.diskSize t.positionMap))

  let flush t =
    let%lwt info = BlockDevice.get_info t.blockDevice in
    let binarySize = bin_size_t t in
    let offset = Int64.(t.offset + diskSize t) in
    let requiredBlocks = (binarySize - 1) / info.BlockDevice.sector_size + 1 in
    let length = requiredBlocks * info.BlockDevice.sector_size in
    let superblock = { offset ; length } in
    let superblockBuffer = Bin_prot.Common.create_buf info.BlockDevice.sector_size in
    let _ = bin_write_superblock superblockBuffer ~pos:0 superblock in
    let superblockCstruct = Cstruct.of_bigarray superblockBuffer in
    BlockDevice.write t.blockDevice 0L [superblockCstruct] >>= fun () ->
    let buffer = Bin_prot.Common.create_buf length in
    let _ = bin_write_t buffer ~pos:0 t in
    let cstruct = Cstruct.of_bigarray buffer in
    BlockDevice.write t.blockDevice offset [cstruct] >>= fun () ->
    return (`Ok ())

  let connect blockDevice =
    let%lwt info = BlockDevice.get_info blockDevice in
    let superblockCstruct = Cstruct.create info.BlockDevice.sector_size in
    BlockDevice.read blockDevice 0L [superblockCstruct] >>= fun () ->
    let superblockBuffer = Cstruct.to_bigarray superblockCstruct in
    let { offset ; length } = bin_read_superblock superblockBuffer ~pos_ref:(ref 0) in
    let core = Cstruct.create length in
    BlockDevice.read blockDevice offset [core] >>= fun () ->
    let buffer = Cstruct.to_bigarray core in
    let t = bin_read_t buffer ~pos_ref:(ref 0) blockDevice in
    return (`Ok t)

  let fakeReconnect oram blockDevice =
    return (`Ok { oram with blockDevice })

  let disconnect t =
    match%lwt flush t with
    | `Ok _ -> BlockDevice.disconnect t.blockDevice
    | `Error _ -> failwith "Failed to flush ORAM! You're probably screwed!"

  let calculateAddressOfBucket t leaf level =
    let rec loop acc = function
      | 0 -> acc
      | y ->
         if Int64.bit_and (Int64.shift_right leaf (y + (t.structuralInfo.height - level) - 1)) 1L = 1L
         then loop Int64.(2L * acc + (t.bucketSize * (of_int t.structuralInfo.sectorsPerBlock) * 2L)) (y - 1)
         else loop Int64.(2L * acc + (t.bucketSize * (of_int t.structuralInfo.sectorsPerBlock))) (y - 1)
    in loop 0L level

  let writeBucket t address bucket =
    (* Printf.printf "%Ld\n" Int64.(address + t.offset); *)
    let buffers = List.map ~f:OBlock.to_cstruct bucket in
    BlockDevice.write t.blockDevice Int64.(address + t.offset) buffers

  let writePathToLeaf t leaf path =
    let rec loop path = function
      | -1 ->
         begin match path with
         | [] -> return (`Ok ())
         | _ -> return (`Error (`Unknown "Wrong number of buckets in path"))
         end
      | y ->
         match path with
         | (bucket :: restOfPath) ->
            let addressOfBucket = calculateAddressOfBucket t leaf (t.structuralInfo.height - y) in
            writeBucket t addressOfBucket bucket >>= fun () ->
            loop restOfPath (y - 1)
         | _ -> return (`Error (`Unknown "Wrong number of buckets in path"))
    in loop path t.structuralInfo.height

  let createBuffersForBucket t =
    let sector_size = t.info.sector_size + 8 in
    let zeroToBucketSize = List.range 0 (Option.value (Int64.to_int t.bucketSize) ~default:4) in
    List.map ~f:(fun _ -> Cstruct.create sector_size) zeroToBucketSize

  let readBucket t address =
    (* Printf.printf "%Ld\n" Int64.(address + t.offset);*)
    let buffersForBucket = createBuffersForBucket t in
    BlockDevice.read t.blockDevice Int64.(address + t.offset) buffersForBucket >>= fun () ->
    let bucket = List.map ~f:OBlock.of_cstruct buffersForBucket in
    return (`Ok bucket)

  let rec readPathToLeafToLevel t leaf = function
    | -1 -> return (`Ok [])
    | level ->
       let addressOfBucket = calculateAddressOfBucket t leaf (level) in
       readBucket t addressOfBucket >>= fun bucket ->
       readPathToLeafToLevel t leaf (level - 1) >>= fun restOfPath ->
       return (`Ok (bucket :: restOfPath))

  let readPathToLeaf t leaf = readPathToLeafToLevel t leaf t.structuralInfo.height

  let validBlocks t l x = List.partition_tf ~f:(fun (pos, _) -> Int64.shift_right pos (t.structuralInfo.height - l) = Int64.shift_right x (t.structuralInfo.height - l))

  let buildBucket t choices =
    let rec loop left acc rest =
      match left, acc, rest with
      | 0L, bs, rest -> (bs, rest)
      | n, acc, [] -> loop Int64.(n - 1L) (OBlock.dummy t.info.sector_size :: acc) []
      | n, acc, ((_, (a,d)) :: xs) -> Stash.remove t.stash a; loop Int64.(n - 1L) ((a,d) :: acc) xs
    in loop t.bucketSize [] choices

  let buildPathToLeaf t leaf =
    let stashList = Stash.to_alist t.stash in
    let rec mapPosition acc = function
      | [] -> return (`Ok acc)
      | ((a,d) :: xs) ->
         PositionMap.get t.positionMap a >>= fun pos ->
         mapPosition ((pos,(a,d)) :: acc) xs
    in
    mapPosition [] stashList >>= fun stash ->
    let rec aux t acc stash x = function
      | 0 ->
         let (hits, misses) = validBlocks t 0 x stash in
         let (bucket, rest) = buildBucket t hits in
         bucket :: acc
      | l ->
         let (hits, misses) = validBlocks t l x stash in
         let (bucket, rest) = buildBucket t hits in
         aux t (bucket :: acc) (rest @ misses) x (l - 1)
    in
    return (`Ok (aux t [] stash leaf t.structuralInfo.height))

  let access t operation address dataToWrite =
    PositionMap.get t.positionMap address >>= fun leaf ->
    let newLeaf = Random.int64 t.structuralInfo.numLeaves in
    PositionMap.set t.positionMap address newLeaf >>= fun () ->
    readPathToLeaf t leaf >>= fun pathRead ->
    List.iter pathRead ~f:(
                List.iter ~f:(
                            fun (address, data) ->  Stash.add t.stash ~address ~data));
    let dataRead = match Stash.find t.stash address with
      | Some data -> data
      | None -> Cstruct.create t.info.sector_size
    in
    begin match operation with
    | Write ->
       begin match dataToWrite with
       | Some data ->
          Stash.add t.stash ~address ~data;
          return (`Ok ())
       | None -> return (`Error (`Unknown "Writing nothing!"))
       end
    | Read -> return (`Ok ())
    end >>= fun () ->
    buildPathToLeaf t leaf >>= fun pathToWrite ->
    writePathToLeaf t leaf pathToWrite >>= fun () ->
    return (`Ok (dataRead))

  let initialise t =
    let dummy_struct = Cstruct.create (t.info.sector_size + 8) in
    Cstruct.LE.set_uint64 dummy_struct 0 (-1L);
    for i = 8 to Cstruct.len dummy_struct - 1 do
      Cstruct.set_uint8 dummy_struct i 0
    done;
    let rec loop = function
      | 0L -> return (`Ok ())
      | x -> BlockDevice.write t.blockDevice Int64.(t.offset + (x - 1L) * (of_int t.structuralInfo.sectorsPerBlock)) [dummy_struct] >>= fun () -> loop Int64.(x - 1L)
    in
    loop t.info.size_sectors >>= fun () ->
    if t.offset > Int64.of_int t.structuralInfo.sectorsPerBlock
    then (
      let b = Float.(to_int (log (of_int t.info.sector_size * 8.) / log 2.)) in
      let height' = t.structuralInfo.height - 6 + b in
      let bound = Int64.(pow 2L (of_int height')) in
      let rec loop = function
        | 0L -> return (`Ok ())
        | x ->
           let data = Cstruct.create t.info.sector_size in
           for i = 0 to (Cstruct.len data / 8) - 1 do
             Cstruct.BE.set_uint64 data (i * 8) (Random.int64 bound);
           done;
           access t Write Int64.(x - 1L) (Some data) >>= fun _ ->
           loop Int64.(x - 1L)
      in
      loop t.info.size_sectors
    ) else return (`Ok ())

  let stashSizesForBucketSize = [(4L,89L);(5L,63L);(6L,53L)]

  let calculateHeightToFillBlockDevice info sectorsPerBlock bucketSize =
    (* Need to accommodate stash space into this calculation *)
    let maxStashSizeInBlocks = List.Assoc.find_exn stashSizesForBucketSize bucketSize in
    let normalisedSizeInSectors = Int64.(info.BlockDevice.size_sectors / of_int sectorsPerBlock) in
    let sectorsForORAM = Int64.(normalisedSizeInSectors - maxStashSizeInBlocks) in
    Printf.printf "Saving %Ld of %Ld blocks for stash, leaving %Ld for ORAM\n%!" maxStashSizeInBlocks normalisedSizeInSectors sectorsForORAM;
    if sectorsForORAM > 4L
    then floor_log Int64.(sectorsForORAM / bucketSize + 1L) - 1
    else failwith "BlockDevice too small for current parameters"

  let calculateHeightOfPositionMap desiredSizeInSectors sector_size bucketSize =
    let reductionFactor = Float.(to_int (log (of_int sector_size * 8.) / log 2.)) in
    floor_log Int64.(desiredSizeInSectors / bucketSize + 1L) + 5 - reductionFactor

  let calculateHeightOfDataORAM desiredSizeInSectors bucketSize =
    floor_log Int64.(desiredSizeInSectors / bucketSize + 1L) - 1

  let calculateHeight desiredSizeInSectors info sectorsPerBlock sector_size bucketSize offset =
    if desiredSizeInSectors = 0L
    then calculateHeightToFillBlockDevice info sectorsPerBlock bucketSize
    else if offset > Int64.of_int (sector_size + sizeOfAddressInBytes)
    then calculateHeightOfPositionMap desiredSizeInSectors sector_size bucketSize
    else calculateHeightOfDataORAM desiredSizeInSectors bucketSize

  let createInstanceOfType desiredSizeInSectors bucketSize desiredBlockSize offset blockDevice =
    let%lwt info = BlockDevice.get_info blockDevice in
    let read_write = info.BlockDevice.read_write in
    let sectorsPerBlock = (desiredBlockSize - 1) / info.BlockDevice.sector_size + 1 in
    let blockOffset = Int64.(offset * (of_int sectorsPerBlock)) in
    let sector_size = sectorsPerBlock * info.BlockDevice.sector_size - sizeOfAddressInBytes in
    (* should in fact save log size_sectors space for the stash *)
    let height = calculateHeight desiredSizeInSectors info sectorsPerBlock sector_size bucketSize blockOffset in
    let numLeaves = Int64.(pow 2L (of_int height)) in
    let size_sectors = Int64.(bucketSize * (2L * numLeaves - 1L)) in
    (* let totalSizeSectors = Int64.(info.BlockDevice.size_sectors / of_int sectorsPerBlock) in *)
    (* Printf.printf "Creating ORAM of size %Ld in block of size %Ld with block offset %Ld\n" size_sectors totalSizeSectors blockOffset; *)
    let stash = Stash.create () in
    PositionMap.create ~desiredSizeInSectors:size_sectors ~bucketSize ~desiredBlockSize ~offset:Int64.(offset + size_sectors) blockDevice >>= fun positionMap ->
    let t = {
        info = { read_write; sector_size; size_sectors } ;
        structuralInfo = { height ; numLeaves ; sectorsPerBlock } ;
        bucketSize ; desiredBlockSize ; offset = blockOffset ;
        stash ; positionMap ; blockDevice
      }
    in
    (*Printf.printf "Created ORAM of size %Ld with block size %d and bucket size %Ld\n%!" size_sectors sector_size bucketSize;*)
    return (`Ok t)

  let create ?(desiredSizeInSectors = 0L) ?(bucketSize = 4L) ?(desiredBlockSize = 0x2000) ?(offset = 1L) blockDevice =
    Random.self_init ();
    let%lwt info = BlockDevice.get_info blockDevice in
    let desiredBlockSizeToUse =
      if info.BlockDevice.sector_size > desiredBlockSize
      then info.BlockDevice.sector_size
      else desiredBlockSize
    in
    (*Printf.printf "Attempting to create ORAM of size %Ld with block size %d and bucket size %Ld\n%!" desiredSizeInSectors desiredBlockSizeToUse bucketSize;*)
    createInstanceOfType desiredSizeInSectors bucketSize desiredBlockSizeToUse offset blockDevice >>= fun t ->
    initialise t >>= fun () ->
    flush t >>= fun () ->
    return (`Ok t)

  let writeBuffer t startAddress buffer =
    let bufferLengthInSectors = buffer.Cstruct.len / t.info.sector_size in
    let rec writeSectors = function
      | -1 -> return (`Ok ())
      | sectorAddressInBuffer ->
         let offsetInBuffer = sectorAddressInBuffer * t.info.sector_size in
         let dataToWrite = Cstruct.sub buffer offsetInBuffer t.info.sector_size in
         access t Write Int64.(startAddress + (of_int sectorAddressInBuffer)) (Some dataToWrite) >>= fun _ ->
         writeSectors (sectorAddressInBuffer - 1)
    in writeSectors (bufferLengthInSectors - 1)
<commit_msg>Added print statements to debug encrypted blocks<commit_after>open Lwt
open Printf
open V1_LWT
open Core_kernel.Std
open PosMapIntf
open Bin_prot.Std

module Make (MakePositionMap : PosMapF) (BlockDevice : BLOCK) = struct

  module PositionMap = MakePositionMap(BlockDevice)

  type 'a io = 'a Lwt.t

  type page_aligned_buffer = Cstruct.t

  type error = [
    | `Unknown of string
    | `Unimplemented
    | `Is_read_only
    | `Disconnected
    ]

  type superblock = {
      offset : int64;
      length : int;
    } [@@ deriving bin_io]

  type info = {
      read_write: bool;
      sector_size: int;
      size_sectors: int64;
    } [@@ deriving bin_io]

  type structuralInfo = {
      height : int;
      numLeaves : int64;
      sectorsPerBlock : int;
    } [@@ deriving bin_io]

  type core = {
      info : info;
      structuralInfo : structuralInfo;
      bucketSize : int64;
      offset : int64;
      desiredBlockSize : int;
      stash : Stash.t;
    } [@@ deriving bin_io]

  type t = {
      info : info;
      structuralInfo : structuralInfo;
      bucketSize : int64;
      offset : int64;
      desiredBlockSize : int;
      stash : Stash.t;
      positionMap : PositionMap.t;
      blockDevice : BlockDevice.t;
    }

  let core { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash } =
    { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash }

  let extendCore ({ info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash } : core) positionMap blockDevice =
    { info ; structuralInfo ; bucketSize ; offset ; desiredBlockSize ; stash ; positionMap ; blockDevice}

  let bin_size_t t =
    let coreSize = core t |> bin_size_core in
    let posMapSize = PositionMap.bin_size_t t.positionMap in
    coreSize + posMapSize

  let bin_write_t buf ~pos t =
    let pos = core t |> bin_write_core buf ~pos in
    PositionMap.bin_write_t buf ~pos t.positionMap

  let bin_writer_t = { Bin_prot.Type_class.size = bin_size_t; write = bin_write_t  }

  let bin_read_t buf ~pos_ref blockDevice =
    let core = bin_read_core buf ~pos_ref in
    let positionMap = PositionMap.bin_read_t buf ~pos_ref blockDevice in
    extendCore core positionMap blockDevice

  type id = string

  type op =
    | Read
    | Write

  type bucket = OBlock.t list

  let sizeOfAddressInBytes = 8

  let ( >>= ) x f = x >>= function
                    | `Error e -> return (`Error e)
                    | `Ok x -> f x

  let get_info { info } = return info

  let getStructuralInfo { structuralInfo } = return structuralInfo

  let floor_log x =
    let rec loop acc = function
      | 1L -> acc
      | x -> loop (acc + 1) Int64.(x / 2L)
    in loop 0 x

  let diskSize t =
    Int64.((t.info.size_sectors * of_int t.structuralInfo.sectorsPerBlock) + (PositionMap.diskSize t.positionMap))

  let flush t =
    let%lwt info = BlockDevice.get_info t.blockDevice in
    let binarySize = bin_size_t t in
    let offset = Int64.(t.offset + diskSize t) in
    let requiredBlocks = (binarySize - 1) / info.BlockDevice.sector_size + 1 in
    let length = requiredBlocks * info.BlockDevice.sector_size in
    let superblock = { offset ; length } in
    let superblockBuffer = Bin_prot.Common.create_buf info.BlockDevice.sector_size in
    let _ = bin_write_superblock superblockBuffer ~pos:0 superblock in
    let superblockCstruct = Cstruct.of_bigarray superblockBuffer in
    BlockDevice.write t.blockDevice 0L [superblockCstruct] >>= fun () ->
    let buffer = Bin_prot.Common.create_buf length in
    let _ = bin_write_t buffer ~pos:0 t in
    let cstruct = Cstruct.of_bigarray buffer in
    BlockDevice.write t.blockDevice offset [cstruct] >>= fun () ->
    return (`Ok ())

  let connect blockDevice =
    let%lwt info = BlockDevice.get_info blockDevice in
    let superblockCstruct = Cstruct.create info.BlockDevice.sector_size in
    BlockDevice.read blockDevice 0L [superblockCstruct] >>= fun () ->
    let superblockBuffer = Cstruct.to_bigarray superblockCstruct in
    let { offset ; length } = bin_read_superblock superblockBuffer ~pos_ref:(ref 0) in
    let core = Cstruct.create length in
    BlockDevice.read blockDevice offset [core] >>= fun () ->
    let buffer = Cstruct.to_bigarray core in
    let t = bin_read_t buffer ~pos_ref:(ref 0) blockDevice in
    return (`Ok t)

  let fakeReconnect oram blockDevice =
    return (`Ok { oram with blockDevice })

  let disconnect t =
    match%lwt flush t with
    | `Ok _ -> BlockDevice.disconnect t.blockDevice
    | `Error _ -> failwith "Failed to flush ORAM! You're probably screwed!"

  let calculateAddressOfBucket t leaf level =
    let rec loop acc = function
      | 0 -> acc
      | y ->
         if Int64.bit_and (Int64.shift_right leaf (y + (t.structuralInfo.height - level) - 1)) 1L = 1L
         then loop Int64.(2L * acc + (t.bucketSize * (of_int t.structuralInfo.sectorsPerBlock) * 2L)) (y - 1)
         else loop Int64.(2L * acc + (t.bucketSize * (of_int t.structuralInfo.sectorsPerBlock))) (y - 1)
    in loop 0L level

  let writeBucket t address bucket =
    (* Printf.printf "%Ld\n" Int64.(address + t.offset); *)
    let buffers = List.map ~f:OBlock.to_cstruct bucket in
    BlockDevice.write t.blockDevice Int64.(address + t.offset) buffers

  let writePathToLeaf t leaf path =
    let rec loop path = function
      | -1 ->
         begin match path with
         | [] -> return (`Ok ())
         | _ -> return (`Error (`Unknown "Wrong number of buckets in path"))
         end
      | y ->
         match path with
         | (bucket :: restOfPath) ->
            let addressOfBucket = calculateAddressOfBucket t leaf (t.structuralInfo.height - y) in
            writeBucket t addressOfBucket bucket >>= fun () ->
            loop restOfPath (y - 1)
         | _ -> return (`Error (`Unknown "Wrong number of buckets in path"))
    in loop path t.structuralInfo.height

  let createBuffersForBucket t =
    let sector_size = t.info.sector_size + 8 in
    let zeroToBucketSize = List.range 0 (Option.value (Int64.to_int t.bucketSize) ~default:4) in
    List.map ~f:(fun _ -> Cstruct.create sector_size) zeroToBucketSize

  let readBucket t address =
    (* Printf.printf "%Ld\n" Int64.(address + t.offset);*)
    let buffersForBucket = createBuffersForBucket t in
    BlockDevice.read t.blockDevice Int64.(address + t.offset) buffersForBucket >>= fun () ->
    let bucket = List.map ~f:OBlock.of_cstruct buffersForBucket in
    return (`Ok bucket)

  let rec readPathToLeafToLevel t leaf = function
    | -1 -> return (`Ok [])
    | level ->
       let addressOfBucket = calculateAddressOfBucket t leaf (level) in
       readBucket t addressOfBucket >>= fun bucket ->
       readPathToLeafToLevel t leaf (level - 1) >>= fun restOfPath ->
       return (`Ok (bucket :: restOfPath))

  let readPathToLeaf t leaf = readPathToLeafToLevel t leaf t.structuralInfo.height

  let validBlocks t l x = List.partition_tf ~f:(fun (pos, _) -> Int64.shift_right pos (t.structuralInfo.height - l) = Int64.shift_right x (t.structuralInfo.height - l))

  let buildBucket t choices =
    let rec loop left acc rest =
      match left, acc, rest with
      | 0L, bs, rest -> (bs, rest)
      | n, acc, [] -> loop Int64.(n - 1L) (OBlock.dummy t.info.sector_size :: acc) []
      | n, acc, ((_, (a,d)) :: xs) -> Stash.remove t.stash a; loop Int64.(n - 1L) ((a,d) :: acc) xs
    in loop t.bucketSize [] choices

  let buildPathToLeaf t leaf =
    let stashList = Stash.to_alist t.stash in
    let rec mapPosition acc = function
      | [] -> return (`Ok acc)
      | ((a,d) :: xs) ->
         PositionMap.get t.positionMap a >>= fun pos ->
         mapPosition ((pos,(a,d)) :: acc) xs
    in
    mapPosition [] stashList >>= fun stash ->
    let rec aux t acc stash x = function
      | 0 ->
         let (hits, misses) = validBlocks t 0 x stash in
         let (bucket, rest) = buildBucket t hits in
         bucket :: acc
      | l ->
         let (hits, misses) = validBlocks t l x stash in
         let (bucket, rest) = buildBucket t hits in
         aux t (bucket :: acc) (rest @ misses) x (l - 1)
    in
    return (`Ok (aux t [] stash leaf t.structuralInfo.height))

  let access t operation address dataToWrite =
    PositionMap.get t.positionMap address >>= fun leaf ->
    let newLeaf = Random.int64 t.structuralInfo.numLeaves in
    PositionMap.set t.positionMap address newLeaf >>= fun () ->
    readPathToLeaf t leaf >>= fun pathRead ->
    List.iter pathRead ~f:(
                List.iter ~f:(
                            fun (address, data) ->  Stash.add t.stash ~address ~data));
    let dataRead = match Stash.find t.stash address with
      | Some data -> data
      | None -> Cstruct.create t.info.sector_size
    in
    begin match operation with
    | Write ->
       begin match dataToWrite with
       | Some data ->
          Stash.add t.stash ~address ~data;
          return (`Ok ())
       | None -> return (`Error (`Unknown "Writing nothing!"))
       end
    | Read -> return (`Ok ())
    end >>= fun () ->
    buildPathToLeaf t leaf >>= fun pathToWrite ->
    writePathToLeaf t leaf pathToWrite >>= fun () ->
    return (`Ok (dataRead))

  let initialise t =
    let dummy_struct = Cstruct.create (t.info.sector_size + 8) in
    Cstruct.LE.set_uint64 dummy_struct 0 (-1L);
    for i = 8 to Cstruct.len dummy_struct - 1 do
      Cstruct.set_uint8 dummy_struct i 0
    done;
    let rec loop = function
      | 0L -> return (`Ok ())
      | x -> BlockDevice.write t.blockDevice Int64.(t.offset + (x - 1L) * (of_int t.structuralInfo.sectorsPerBlock)) [dummy_struct] >>= fun () -> loop Int64.(x - 1L)
    in
    loop t.info.size_sectors >>= fun () ->
    if t.offset > Int64.of_int t.structuralInfo.sectorsPerBlock
    then (
      let b = Float.(to_int (log (of_int t.info.sector_size * 8.) / log 2.)) in
      let height' = t.structuralInfo.height - 6 + b in
      let bound = Int64.(pow 2L (of_int height')) in
      let rec loop = function
        | 0L -> return (`Ok ())
        | x ->
           let data = Cstruct.create t.info.sector_size in
           for i = 0 to (Cstruct.len data / 8) - 1 do
             Cstruct.BE.set_uint64 data (i * 8) (Random.int64 bound);
           done;
           access t Write Int64.(x - 1L) (Some data) >>= fun _ ->
           loop Int64.(x - 1L)
      in
      loop t.info.size_sectors
    ) else return (`Ok ())

  let stashSizesForBucketSize = [(4L,89L);(5L,63L);(6L,53L)]

  let calculateHeightToFillBlockDevice info sectorsPerBlock bucketSize =
    (* Need to accommodate stash space into this calculation *)
    let maxStashSizeInBlocks = List.Assoc.find_exn stashSizesForBucketSize bucketSize in
    let normalisedSizeInSectors = Int64.(info.BlockDevice.size_sectors / of_int sectorsPerBlock) in
    let sectorsForORAM = Int64.(normalisedSizeInSectors - maxStashSizeInBlocks) in
    Printf.printf "Saving %Ld of %Ld blocks for stash, leaving %Ld for ORAM\n%!" maxStashSizeInBlocks normalisedSizeInSectors sectorsForORAM;
    if sectorsForORAM > 4L
    then floor_log Int64.(sectorsForORAM / bucketSize + 1L) - 1
    else failwith "BlockDevice too small for current parameters"

  let calculateHeightOfPositionMap desiredSizeInSectors sector_size bucketSize =
    let reductionFactor = Float.(to_int (log (of_int sector_size * 8.) / log 2.)) in
    floor_log Int64.(desiredSizeInSectors / bucketSize + 1L) + 5 - reductionFactor

  let calculateHeightOfDataORAM desiredSizeInSectors bucketSize =
    floor_log Int64.(desiredSizeInSectors / bucketSize + 1L) - 1

  let calculateHeight desiredSizeInSectors info sectorsPerBlock sector_size bucketSize offset =
    if desiredSizeInSectors = 0L
    then calculateHeightToFillBlockDevice info sectorsPerBlock bucketSize
    else if offset > Int64.of_int (sector_size + sizeOfAddressInBytes)
    then calculateHeightOfPositionMap desiredSizeInSectors sector_size bucketSize
    else calculateHeightOfDataORAM desiredSizeInSectors bucketSize

  let createInstanceOfType desiredSizeInSectors bucketSize desiredBlockSize offset blockDevice =
    let%lwt info = BlockDevice.get_info blockDevice in
    let read_write = info.BlockDevice.read_write in
    let sectorsPerBlock = (desiredBlockSize - 1) / info.BlockDevice.sector_size + 1 in
    let blockOffset = Int64.(offset * (of_int sectorsPerBlock)) in
    let sector_size = sectorsPerBlock * info.BlockDevice.sector_size - sizeOfAddressInBytes in
    (* should in fact save log size_sectors space for the stash *)
    let height = calculateHeight desiredSizeInSectors info sectorsPerBlock sector_size bucketSize blockOffset in
    let numLeaves = Int64.(pow 2L (of_int height)) in
    let size_sectors = Int64.(bucketSize * (2L * numLeaves - 1L)) in
    (* let totalSizeSectors = Int64.(info.BlockDevice.size_sectors / of_int sectorsPerBlock) in *)
    (* Printf.printf "Creating ORAM of size %Ld in block of size %Ld with block offset %Ld\n" size_sectors totalSizeSectors blockOffset; *)
    let stash = Stash.create () in
    PositionMap.create ~desiredSizeInSectors:size_sectors ~bucketSize ~desiredBlockSize ~offset:Int64.(offset + size_sectors) blockDevice >>= fun positionMap ->
    let t = {
        info = { read_write; sector_size; size_sectors } ;
        structuralInfo = { height ; numLeaves ; sectorsPerBlock } ;
        bucketSize ; desiredBlockSize ; offset = blockOffset ;
        stash ; positionMap ; blockDevice
      }
    in
    Printf.printf "Created ORAM of size %Ld with block size %d\n%!" size_sectors sector_size;
    return (`Ok t)

  let create ?(desiredSizeInSectors = 0L) ?(bucketSize = 4L) ?(desiredBlockSize = 0x2000) ?(offset = 1L) blockDevice =
    Random.self_init ();
    let%lwt info = BlockDevice.get_info blockDevice in
    let desiredBlockSizeToUse =
      if info.BlockDevice.sector_size > desiredBlockSize
      then info.BlockDevice.sector_size
      else desiredBlockSize
    in
    (*Printf.printf "Attempting to create ORAM of size %Ld with block size %d and bucket size %Ld\n%!" desiredSizeInSectors desiredBlockSizeToUse bucketSize;*)
    createInstanceOfType desiredSizeInSectors bucketSize desiredBlockSizeToUse offset blockDevice >>= fun t ->
    initialise t >>= fun () ->
    flush t >>= fun () ->
    return (`Ok t)

  let writeBuffer t startAddress buffer =
    let bufferLengthInSectors = buffer.Cstruct.len / t.info.sector_size in
    let rec writeSectors = function
      | -1 -> return (`Ok ())
      | sectorAddressInBuffer ->
         let offsetInBuffer = sectorAddressInBuffer * t.info.sector_size in
         let dataToWrite = Cstruct.sub buffer offsetInBuffer t.info.sector_size in
         access t Write Int64.(startAddress + (of_int sectorAddressInBuffer)) (Some dataToWrite) >>= fun _ ->
         writeSectors (sectorAddressInBuffer - 1)
    in writeSectors (bufferLengthInSectors - 1)

{'fraction_non_alphanumeric': 0.0957023343678316, 'fraction_numerical': 0.012346511941708271, 'mean_word_length': 2.878205128205128, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 243, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '285512', 'n_tokens_mistral': 9675, 'n_tokens_neox': 9174, 'n_words': 2795}