free.hs

{-# LANGUAGE GADTs, FlexibleInstances, GeneralizedNewtypeDeriving, DeriveFunctor, DeriveFoldable, DeriveTraversable, TemplateHaskell, FlexibleContexts, OverloadedStrings, ViewPatterns, RankNTypes, NoMonomorphismRestriction #-}

module Main where

  import Prelude hiding (break)
  import Control.Lens
  import Control.Monad
  import Control.Monad.Free
  import Control.Monad.Free.TH
  --import Data.ByteString (ByteString)
  import Data.Bifunctor
  import Data.Foldable (Foldable)
  import Data.Traversable (Traversable, traverse)
  import Data.String
  --import Text.PrettyPrint (Doc)
  import Text.PrettyPrint.Free

  data IterKind 
    = While
    | Switch
    | For
    deriving (Show, Eq)

  instance Pretty IterKind where
    pretty While = "while"
    pretty Switch = "switch"
    pretty For = "for"

  data JumpKind b
   = Break
   | Continue
   | Returning (Maybe b)

  instance (Pretty b) => Pretty (JumpKind b) where
    pretty Break = "break"
    pretty Continue = "continue"
    pretty (Returning Nothing) = "return"
    pretty (Returning (Just a)) = "return" <+> pretty a

  data LabelKind b
    = Custom String
    | Default
    | Case b

  instance (Pretty b) => Pretty (LabelKind b) where
    pretty (Custom s) = pretty s <> ":"
    pretty Default = "default:"
    pretty (Case b) = "case" <+> pretty b <> ":"

  data Literal 
    = IntLit Integer
    | FltLit Float
    | StrLit String
    | ChrLit Char
    | Ident String

  instance Pretty Literal where
    pretty (IntLit i) = pretty i
    pretty (FltLit f) = pretty f
    pretty (StrLit s) = dquotes $ pretty s
    pretty (ChrLit c) = squotes $ pretty c
    pretty (Ident  i) = pretty i

  data Stmt exp next
    = Compound  { _body :: Statement () 
                , _next :: next 
                }
    | Empty     { _next :: next }
    | IfThen    { _expr :: exp
                , _body :: Statement () 
                , _other :: Maybe (Statement exp)
                , _next :: next 
                }
    | Iteration { _iterKind :: IterKind
                , _expr :: exp
                , _body :: Statement () 
                , _next :: next
                }
    | Jump      { _jumpKind :: JumpKind exp }
    | Labelled  { _labelKind :: LabelKind exp
                , _body :: Statement () 
                , _next :: next
                }
    deriving (Functor, Foldable, Traversable)

  type Statement = Free (Stmt Literal)

  
  makeFree ''Stmt
  makeLenses ''Stmt

  return_ :: MonadFree (Stmt exp) m => m ()
  return_ = jump $ Returning Nothing

  return' :: MonadFree (Stmt s) m => s -> m s
  return' a = jump $ Returning $ Just a

  class (Functor f) => PrettyAlg f where
    prettyA :: f (Doc e) -> (Doc e)

  free :: (Functor f) => (f b -> b) -> (a -> b) -> Free f a -> b
  free _ g (Pure x) = g x
  free f g (Free t) = f (fmap (free f g) t)

  instance (PrettyAlg f, Pretty a) => Pretty (Free f a) where
    pretty = free prettyA (const "") 

  --instance Pretty (Statement b) where
  --  pretty = iter prettyA

  instance (Pretty b) => PrettyAlg (Stmt b) where
    prettyA (Compound b n) = "{" `above` (indent 4 (pretty b)) `above` "}" `above` n
    prettyA (Empty n)      = ";" `above` n
    prettyA (IfThen e b Nothing n) = "if" <+> parens (pretty e) <+> pretty b `above` n
    prettyA (IfThen e b (Just o) n) = "if" <+> parens (pretty e) <+> pretty b `above` "else" <+> pretty o `above` n
    prettyA (Iteration k e b n) = pretty k <+> parens (pretty e) <+> pretty b `above` n
    prettyA (Jump j) = pretty j
    prettyA (Labelled l b n) = pretty l <+> pretty b `above` n

  -- exp :: Traversal (Stmt b next) (Stmt b' next) b b'
  -- next :: Traversal (Stmt b next) (Stmt b next') next 'next

  makePrisms ''Stmt
  -- _Output :: Prism' (Stmt b next) (Stmt b' next) (b, next) (b', next)
  -- _Bell :: Prism' (Stmt b next) next
  -- _Break :: Prism' (Stmt b next) ()

  --output :: MonadFree (Stmt b0) m0 => b0 -> m0 ()
  --output = liftM ()
  -- bell :: MonadFree (Stmt b0) m0 => m0 ()
  -- break :: MonadFree (Stmt b0) m0 => m0 a0

  compound' = void . compound

  sampleProgram :: Statement ()
  sampleProgram = compound $ void $ do
    empty
    iteration While (IntLit 1) $ compound $ void $ do
      empty
      empty
      return' (FltLit 10)
    return' (IntLit 5)

  --instance (Pretty a) => Pretty (Free (Stmt a) ()) where
  --  pretty = iter prettyA

  main = print $ pretty sampleProgram