{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Compiling abstracted combinator expressions ('BTerm') to native
-- Haskell terms.  This can supposedly be more efficient than directly
-- interpreting 'BTerm's, but some benchmarking is probably needed to
-- decide whether we want this or not.
--
-- For more info, see:
--
--   https://byorgey.wordpress.com/2023/07/13/compiling-to-intrinsically-typed-combinators/
module Swarm.Game.World.Compile where

import Data.ByteString (ByteString)
import Data.Hash.Murmur (murmur3)
import Data.Kind (Constraint)
import Data.Tagged (Tagged (unTagged))
import Numeric.Noise.Perlin (noiseValue, perlin)
import Swarm.Game.Location (pattern Location)
import Swarm.Game.World.Abstract (BTerm (..))
import Swarm.Game.World.Coords (Coords (..), coordsToLoc)
import Swarm.Game.World.Gen (Seed)
import Swarm.Game.World.Interpret (interpReflect, interpRot)
import Swarm.Game.World.Syntax (Axis (..), Rot, World)
import Swarm.Game.World.Typecheck (Applicable (..), Const (..), Empty (..), NotFun, Over (..))
import Witch (from)
import Witch.Encoding qualified as Encoding

data CTerm a where
  CFun :: (CTerm a -> CTerm b) -> CTerm (a -> b)
  CConst :: (NotFun a) => a -> CTerm a

instance Applicable CTerm where
  CFun CTerm a -> CTerm b
f $$ :: forall a b. CTerm (a -> b) -> CTerm a -> CTerm b
$$ CTerm a
x = CTerm a -> CTerm b
f CTerm a
x

compile :: Seed -> BTerm a -> CTerm a
compile :: forall a. Seed -> BTerm a -> CTerm a
compile Seed
seed (BApp BTerm (a1 -> a)
b1 BTerm a1
b2) = forall a. Seed -> BTerm a -> CTerm a
compile Seed
seed BTerm (a1 -> a)
b1 forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ forall a. Seed -> BTerm a -> CTerm a
compile Seed
seed BTerm a1
b2
compile Seed
seed (BConst Const a
c) = forall a. Seed -> Const a -> CTerm a
compileConst Seed
seed Const a
c

compileConst :: Seed -> Const a -> CTerm a
compileConst :: forall a. Seed -> Const a -> CTerm a
compileConst Seed
seed = \case
  CLit a
i -> forall a. NotFun a => a -> CTerm a
CConst a
i
  CCell CellVal
c -> forall a. NotFun a => a -> CTerm a
CConst CellVal
c
  Const a
CFI -> forall a b. (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary forall a b. (Integral a, Num b) => a -> b
fromIntegral
  Const a
CIf -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Bool
b) -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
t -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
e -> if Bool
b then CTerm a1
t else CTerm a1
e
  Const a
CNot -> forall a b. (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary Bool -> Bool
not
  Const a
CNeg -> forall a b. (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary forall a. Num a => a -> a
negate
  Const a
CAbs -> forall a b. (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary forall a. Num a => a -> a
abs
  Const a
CAnd -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary Bool -> Bool -> Bool
(&&)
  Const a
COr -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary Bool -> Bool -> Bool
(||)
  Const a
CAdd -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Num a => a -> a -> a
(+)
  Const a
CSub -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary (-)
  Const a
CMul -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Num a => a -> a -> a
(*)
  Const a
CDiv -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Fractional a => a -> a -> a
(/)
  Const a
CIDiv -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Integral a => a -> a -> a
div
  Const a
CMod -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Integral a => a -> a -> a
mod
  Const a
CEq -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Eq a => a -> a -> Bool
(==)
  Const a
CNeq -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Eq a => a -> a -> Bool
(/=)
  Const a
CLt -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Ord a => a -> a -> Bool
(<)
  Const a
CLeq -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Ord a => a -> a -> Bool
(<=)
  Const a
CGt -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Ord a => a -> a -> Bool
(>)
  Const a
CGeq -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall a. Ord a => a -> a -> Bool
(>=)
  Const a
CMask -> forall a.
(NotFun a, Empty a) =>
CTerm (World Bool -> World a -> World a)
compileMask
  Const a
CSeed -> forall a. NotFun a => a -> CTerm a
CConst (forall a b. (Integral a, Num b) => a -> b
fromIntegral Seed
seed)
  CCoord Axis
ax -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst (Coords -> Location
coordsToLoc -> Location Int32
x Int32
y)) -> forall a. NotFun a => a -> CTerm a
CConst (forall a b. (Integral a, Num b) => a -> b
fromIntegral (case Axis
ax of Axis
X -> Int32
x; Axis
Y -> Int32
y))
  Const a
CHash -> CTerm (Coords -> Integer)
compileHash
  Const a
CPerlin -> CTerm (Integer -> Integer -> Double -> Double -> World Double)
compilePerlin
  CReflect Axis
ax -> forall a. Axis -> CTerm (World a -> World a)
compileReflect Axis
ax
  CRot Rot
rot -> forall a. Rot -> CTerm (World a -> World a)
compileRot Rot
rot
  Const a
COver -> forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary forall m. Over m => m -> m -> m
(<!>)
  Const a
K -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
x -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const CTerm a1
x
  Const a
S -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (a1 -> b -> c)
f -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (a1 -> b)
g -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
x -> CTerm (a1 -> b -> c)
f forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm a1
x forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ (CTerm (a1 -> b)
g forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm a1
x)
  Const a
I -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a. a -> a
id
  Const a
B -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (b -> c)
f -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (a1 -> b)
g -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
x -> CTerm (b -> c)
f forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ (CTerm (a1 -> b)
g forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm a1
x)
  Const a
C -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (a1 -> b -> c)
f -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm b
x -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm a1
y -> CTerm (a1 -> b -> c)
f forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm a1
y forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm b
x
  Const a
Φ -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (a1 -> b -> c)
c -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (d -> a1)
f -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (d -> b)
g -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm d
x -> CTerm (a1 -> b -> c)
c forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ (CTerm (d -> a1)
f forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm d
x) forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ (CTerm (d -> b)
g forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm d
x)

unary :: (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary :: forall a b. (NotFun a, NotFun b) => (a -> b) -> CTerm (a -> b)
unary a -> b
op = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst a
x) -> forall a. NotFun a => a -> CTerm a
CConst (a -> b
op a
x)

binary :: (NotFun a, NotFun b, NotFun c) => (a -> b -> c) -> CTerm (a -> b -> c)
binary :: forall a b c.
(NotFun a, NotFun b, NotFun c) =>
(a -> b -> c) -> CTerm (a -> b -> c)
binary a -> b -> c
op = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst a
x) -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst b
y) -> forall a. NotFun a => a -> CTerm a
CConst (a -> b -> c
op a
x b
y)

-- Note we could desugar 'mask p a -> if p a empty' but that would
-- require an explicit 'empty' node, whose type can't be inferred.
compileMask :: (NotFun a, Empty a) => CTerm (World Bool -> World a -> World a)
compileMask :: forall a.
(NotFun a, Empty a) =>
CTerm (World Bool -> World a -> World a)
compileMask = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (World Bool)
p -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (World a)
a -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm Coords
ix ->
  case CTerm (World Bool)
p forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm Coords
ix of
    CConst Bool
b -> if Bool
b then CTerm (World a)
a forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ CTerm Coords
ix else forall a. NotFun a => a -> CTerm a
CConst forall e. Empty e => e
empty

compileHash :: CTerm (Coords -> Integer)
compileHash :: CTerm (Coords -> Integer)
compileHash = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst (Coords (Int32, Int32)
ix)) -> forall a. NotFun a => a -> CTerm a
CConst (forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int32, Int32) -> Word32
h (Int32, Int32)
ix))
 where
  h :: (Int32, Int32) -> Word32
h = Word32 -> ByteString -> Word32
murmur3 Word32
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (s :: k) b. Tagged s b -> b
unTagged forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall source target. From source target => source -> target
from @String @(Encoding.UTF_8 ByteString) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show

compilePerlin :: CTerm (Integer -> Integer -> Double -> Double -> World Double)
compilePerlin :: CTerm (Integer -> Integer -> Double -> Double -> World Double)
compilePerlin =
  forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Integer
s) ->
    forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Integer
o) ->
      forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Double
k) ->
        forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Double
p) ->
          let noise :: Perlin
noise = Seed -> Seed -> Double -> Double -> Perlin
perlin (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
s) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
o) Double
k Double
p
           in forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst (Coords (Int32, Int32)
ix)) -> forall a. NotFun a => a -> CTerm a
CConst (forall {a} {a} {a}.
(Noise a, Integral a, Integral a) =>
(a, a) -> a -> Double
sample (Int32, Int32)
ix Perlin
noise)
 where
  sample :: (a, a) -> a -> Double
sample (a
i, a
j) a
noise = forall a. Noise a => a -> Point -> Double
noiseValue a
noise (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
i forall a. Fractional a => a -> a -> a
/ Double
2, forall a b. (Integral a, Num b) => a -> b
fromIntegral a
j forall a. Fractional a => a -> a -> a
/ Double
2, Double
0)

compileReflect :: Axis -> CTerm (World a -> World a)
compileReflect :: forall a. Axis -> CTerm (World a -> World a)
compileReflect Axis
ax = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (World a)
w -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Coords
c) -> CTerm (World a)
w forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ forall a. NotFun a => a -> CTerm a
CConst (Axis -> Coords -> Coords
interpReflect Axis
ax Coords
c)

compileRot :: Rot -> CTerm (World a -> World a)
compileRot :: forall a. Rot -> CTerm (World a -> World a)
compileRot Rot
rot = forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \CTerm (World a)
w -> forall a b. (CTerm a -> CTerm b) -> CTerm (a -> b)
CFun forall a b. (a -> b) -> a -> b
$ \(CConst Coords
c) -> CTerm (World a)
w forall (t :: * -> *) a b. Applicable t => t (a -> b) -> t a -> t b
$$ forall a. NotFun a => a -> CTerm a
CConst (Rot -> Coords -> Coords
interpRot Rot
rot Coords
c)

type family NoFunParams a :: Constraint where
  NoFunParams (a -> b) = (NotFun a, NoFunParams b)
  NoFunParams _ = ()

-- | Interpret a compiled term into the host language.
runCTerm :: (NoFunParams a) => CTerm a -> a
runCTerm :: forall a. NoFunParams a => CTerm a -> a
runCTerm (CConst a
a) = a
a
runCTerm (CFun CTerm a -> CTerm b
f) = forall a. NoFunParams a => CTerm a -> a
runCTerm forall b c a. (b -> c) -> (a -> b) -> a -> c
. CTerm a -> CTerm b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NotFun a => a -> CTerm a
CConst