Combinators.hs

{-# LANGUAGE AllowAmbiguousTypes, UndecidableInstances #-}
-- | Iteration combinators over generated Knightian names.
module Imp.DSL.Combinators
  ( GenNames
  , ConcatMapTag
  , MapName(..)
  , intervalMap
  , knightMap
  , intervalN
  , knightN
  , TagFoldM(..)
  , foldMN
  , tagFold
  , tagMap
  , tagScan
  , tagN
  , foldN
  , scanN
  ) where

import GHC.TypeLits
  ( Natural, Symbol, AppendSymbol, ConsSymbol, NatToChar, KnownSymbol
  , Div, Mod, type (+), type (-) )

import Prelude hiding (return, (>>=), (>>), flip, fmap)
import Data.Proxy (Proxy(..))
import Imp.DSL.Grade (Merge, TagAll)
import Imp.DSL (Imp(..), knight, interval)

-- | Convert a type-level 'Natural' to its decimal 'Symbol' representation.
type family NatToSymbol (n :: Natural) :: Symbol where
  NatToSymbol n = NatToSymbolGo (Div n 10) (Mod n 10)

type family NatToSymbolGo (q :: Natural) (r :: Natural) :: Symbol where
  NatToSymbolGo 0 r = ConsSymbol (NatToChar (r + 48)) ""
  NatToSymbolGo q r = AppendSymbol (NatToSymbol q) (ConsSymbol (NatToChar (r + 48)) "")

-- | Number of decimal digits in a 'Natural'.
type family Digits (n :: Natural) :: Natural where
  Digits n = DigitsH (Div n 10)

type family DigitsH (q :: Natural) :: Natural where
  DigitsH 0 = 1
  DigitsH q = 1 + DigitsH (Div q 10)

-- | @k@ zeros: @Zeros 3 = \"000\"@.
type family Zeros (k :: Natural) :: Symbol where
  Zeros 0 = ""
  Zeros k = AppendSymbol "0" (Zeros (k - 1))

-- | Decimal representation of @n@, left-padded with zeros to @width@ digits.
type PadNat (width :: Natural) (n :: Natural) =
  AppendSymbol (Zeros (width - Digits n)) (NatToSymbol n)

-- | Generate a list of 1-indexed numbered names.
--   Padding matters because grades are kept sorted lexicographically.
type family GenNames (count :: Natural) (base :: Symbol) :: [Symbol] where
  GenNames 0 _    = '[]
  GenNames n base = GenNamesGo n (Digits n) base

type family GenNamesGo (n :: Natural) (width :: Natural) (base :: Symbol) :: [Symbol] where
  GenNamesGo 0 _     _    = '[]
  GenNamesGo n width base =
    Merge (GenNamesGo (n - 1) width base) '[AppendSymbol base (PadNat width n)]

-- | Map 'TagAll' over a list of tags and concatenate the results.
type family ConcatMapTag (tags :: [Symbol]) (k :: [Symbol]) :: [Symbol] where
  ConcatMapTag '[]       _ = '[]
  ConcatMapTag (t ': ts) k = Merge (TagAll t k) (ConcatMapTag ts k)

-- | Graded @map@: traverse a type-level list of Knightian names,
--   applying a single-name computation to each.
--
--   @
--   mapName \@'[\"x\", \"y\"] knight   -- two independent Knightian choices
--   @
class MapName (names :: [Symbol]) where
  mapName :: Ord a => (forall n. KnownSymbol n => Imp '[n] a) -> Imp names [a]

instance MapName '[] where
  mapName _ = ImpReturn []

instance (KnownSymbol n, MapName ns, Merge '[n] ns ~ (n ': ns)) => MapName (n ': ns) where
  mapName f = ImpBind (f @n) $ \x ->
    ImpBind (mapName @ns f) (ImpReturn . (x :))

-- | Independent Knightian choice per name.
knightMap :: MapName names => Imp names [Bool]
knightMap = mapName knight

-- | Independent interval per name.
intervalMap :: MapName names => Double -> Double -> Imp names [Bool]
intervalMap lo hi = mapName (interval lo hi)

-- | Independent numbered Knightian choices.
knightN :: forall n base.
  MapName (GenNames n base) =>
  Imp (GenNames n base) [Bool]
knightN = knightMap @(GenNames n base)

-- | Independent numbered intervals.
intervalN :: forall n base.
  MapName (GenNames n base) =>
  Double -> Double -> Imp (GenNames n base) [Bool]
intervalN lo hi = intervalMap @(GenNames n base) lo hi

-- | Graded left fold where the step function is monadic with unit grade.
--   Enables conditioning/flips at each time step.
--
--   @
--   tagFoldM \@'[\"move1\", \"move2\"] robotDynamics P1 $ \\pos move -> Imp.do
--     let pos' = step3 pos move
--     observe (pos' /= P0)
--     Imp.return pos'
--   @
class TagFoldM (tags :: [Symbol]) (g :: [Symbol]) where
  tagFoldM :: (Ord a, Ord b)
           => Imp g a -> b -> (b -> a -> Imp '[] b) -> Imp (ConcatMapTag tags g) b

instance TagFoldM '[] g where
  tagFoldM _ acc _ = ImpReturn acc

instance ( KnownSymbol t
         , TagFoldM ts g
         ) => TagFoldM (t ': ts) g where
  tagFoldM prog acc f = ImpBind (ImpTag (Proxy @t) prog) $ \x ->
    ImpBind (f acc x) $ \acc' ->
      tagFoldM @ts @g prog acc' f

-- | @foldMN \@n \@base prog acc f@: numbered monadic fold.
foldMN :: forall n base g a b.
  (TagFoldM (GenNames n base) g, Ord a, Ord b) =>
  Imp g a -> b -> (b -> a -> Imp '[] b) -> Imp (ConcatMapTag (GenNames n base) g) b
foldMN = tagFoldM @(GenNames n base) @g

-- | Graded left fold over tagged subprograms.
--
--   @
--   tagFold \@'[\"move1\", \"move2\"] robotDynamics P0 step3
--   @
tagFold :: forall tags g a b. (TagFoldM tags g, Ord a, Ord b) =>
  Imp g a -> b -> (b -> a -> b) -> Imp (ConcatMapTag tags g) b
tagFold prog acc f = tagFoldM @tags prog acc (\b a -> ImpReturn (f b a))

-- | Graded @map@ over tagged subprograms.
--
--   @
--   tagMap \@'[\"move1\", \"move2\", \"move3\"] robotDynamics
--   @
tagMap :: forall tags g a. (TagFoldM tags g, Ord a)
       => Imp g a -> Imp (ConcatMapTag tags g) [a]
tagMap prog = ImpBind (tagFold @tags prog [] (\as a -> a : as)) (ImpReturn . reverse)

-- | Graded left scan over tagged subprograms.
--
--   @
--   tagScan \@'[\"move1\", \"move2\", \"move3\"] robotDynamics P1 step3
--   @
tagScan :: forall tags g a b. (TagFoldM tags g, Ord a, Ord b) =>
  Imp g a -> b -> (b -> a -> b) -> Imp (ConcatMapTag tags g) [b]
tagScan prog acc f =
  ImpBind (tagFold @tags prog (acc, []) (\(b, bs) a -> let b' = f b a in (b', b' : bs)))
          (ImpReturn . reverse . snd)

-- | @tagN \@n \@base prog@: numbered tag iteration.
tagN :: forall n base g a.
  (TagFoldM (GenNames n base) g, Ord a) =>
  Imp g a -> Imp (ConcatMapTag (GenNames n base) g) [a]
tagN = tagMap @(GenNames n base)

-- | @foldN \@n \@base prog acc f@: numbered fold over tagged subprograms.
foldN :: forall n base g a b.
  (TagFoldM (GenNames n base) g, Ord a, Ord b) =>
  Imp g a -> b -> (b -> a -> b) -> Imp (ConcatMapTag (GenNames n base) g) b
foldN = tagFold @(GenNames n base) @g

-- | @scanN \@n \@base prog acc f@: numbered scan over tagged subprograms.
scanN :: forall n base g a b.
  (TagFoldM (GenNames n base) g, Ord a, Ord b) =>
  Imp g a -> b -> (b -> a -> b) -> Imp (ConcatMapTag (GenNames n base) g) [b]
scanN = tagScan @(GenNames n base) @g