-- Copyright 2020 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Dex.Foreign.Serialize (
  CAtom,
  dexPrint, dexToCAtom, dexFromCAtom
  ) where

import Control.Monad.IO.Class
import Data.Word
import Data.Functor

import Foreign.C
import Foreign.Ptr
import Foreign.Storable

import IRVariants
import Name
import TopLevel
import Types.Core hiding (CAtom)
import Types.Primitives

import Dex.Foreign.Context
import Dex.Foreign.Util

-- TODO: Free!
dexPrint :: Ptr Context -> Ptr AtomEx -> IO CString
dexPrint contextPtr atomPtr = do
  AtomEx atom <- fromStablePtr atomPtr
  runTopperMFromContext contextPtr do
    -- TODO: Check consistency of atom and context
    liftIO . newCString =<< printCodegen (unsafeCoerceE atom)

data CAtom = CLit LitVal | CRectArray (Ptr ()) [Int] [Int]

instance Storable CAtom where
  sizeOf _ = tag + val + val + val
    where tag = 8; val = 8
  alignment _ = 8
  peek addr = do
    tag <- val @Word64 0
    case tag of
      0 -> do
        litTag <- val @Word64 1
        CLit <$> case litTag of
                   0 -> Int64Lit   <$> val 2
                   1 -> Int32Lit   <$> val 2
                   2 -> Word8Lit   <$> val 2
                   3 -> Float64Lit <$> val 2
                   4 -> Float32Lit <$> val 2
                   5 -> Word32Lit  <$> val 2
                   6 -> Word64Lit  <$> val 2
                   _ -> error "Invalid tag"
      _ -> error "Invalid tag"
    where
      val :: forall a. Storable a => Int -> IO a
      val i = peekByteOff (castPtr addr) (i * 8)
  poke addr catom = case catom of
    CLit lit -> do
      val @Word64 0 0
      case lit of
        Int64Lit   v -> val @Word64 1 0 >> val 2 v
        Int32Lit   v -> val @Word64 1 1 >> val 2 v
        Word8Lit   v -> val @Word64 1 2 >> val 2 v
        Float64Lit v -> val @Word64 1 3 >> val 2 v
        Float32Lit v -> val @Word64 1 4 >> val 2 v
        Word32Lit  v -> val @Word64 1 5 >> val 2 v
        Word64Lit  v -> val @Word64 1 6 >> val 2 v
        PtrLit   _ _ -> error "Unsupported"
    CRectArray _ _ _ -> error "Unsupported"
    where
      val :: forall a. Storable a => Int -> a -> IO ()
      val i v = pokeByteOff (castPtr addr) (i * 8) v

dexToCAtom :: Ptr AtomEx -> Ptr CAtom -> IO CInt
dexToCAtom atomPtr resultPtr = do
  AtomEx atom <- fromStablePtr atomPtr
  scalarAtomToCAtom atom
  where
    scalarAtomToCAtom :: Atom CoreIR n -> IO CInt
    scalarAtomToCAtom atom = case atom of
      Con con -> case con of
        Lit l -> poke resultPtr (CLit l) $> 1
        _ -> notSerializable
      NewtypeCon NatCon rep -> scalarAtomToCAtom rep
      _ -> notSerializable

    notSerializable = setError "Unserializable atom" $> 0

dexFromCAtom :: Ptr CAtom -> IO (Ptr AtomEx)
dexFromCAtom catomPtr = do
  catom <- peek catomPtr
  case catom of
    CLit lit         -> toStablePtr $ AtomEx $ Con $ Lit lit
    CRectArray _ _ _ -> unsupported
  where
    unsupported = setError "Unsupported CAtom" $> nullPtr