disco-0.2: src/Disco/AST/Core.hs
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-}
-- |
-- Module : Disco.AST.Core
-- Copyright : disco team and contributors
-- Maintainer : byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Abstract syntax trees representing the desugared, untyped core
-- language for Disco.
module Disco.AST.Core (
-- * Core AST
ShouldMemo (..),
Core (..),
Op (..),
opArity,
substQC,
substsQC,
)
where
import Control.Lens.Plated
import Data.Data (Data)
import Data.Data.Lens (uniplate)
import Data.Ratio
import qualified Data.Set as S
import Disco.AST.Generic (Side, selectSide)
import Disco.Effects.LFresh
import Disco.Names (QName)
import Disco.Pretty
import Disco.Syntax.Operators (BOp (..))
import Disco.Types
import GHC.Generics
import Polysemy (Members, Sem)
import Polysemy.Reader
import Unbound.Generics.LocallyNameless hiding (LFresh, lunbind)
import Prelude hiding ((<>))
import qualified Prelude as P
data ShouldMemo = Memo | NoMemo deriving (Show, Generic, Data, Alpha)
-- | AST for the desugared, untyped core language.
data Core where
-- | A variable.
CVar :: QName Core -> Core
-- | A rational number.
CNum :: Rational -> Core
-- | A built-in constant.
CConst :: Op -> Core
-- | An injection into a sum type, i.e. a value together with a tag
-- indicating which element of a sum type we are in. For example,
-- false is represented by @CSum L CUnit@; @right(v)@ is
-- represented by @CSum R v@. Note we do not need to remember
-- which type the constructor came from; if the program
-- typechecked then we will never end up comparing constructors
-- from different types.
CInj :: Side -> Core -> Core
-- | A primitive case expression on a value of a sum type.
CCase :: Core -> Bind (Name Core) Core -> Bind (Name Core) Core -> Core
-- | The unit value.
CUnit :: Core
-- | A pair of values.
CPair :: Core -> Core -> Core
-- | A projection from a product type, i.e. @fst@ or @snd@.
CProj :: Side -> Core -> Core
-- | An anonymous function.
CAbs :: ShouldMemo -> Bind [Name Core] Core -> Core
-- | Function application.
CApp :: Core -> Core -> Core
-- | A "test frame" under which a test case is run. Records the
-- types and legible names of the variables that should
-- be reported to the user if the test fails.
CTest :: [(String, Type, Name Core)] -> Core -> Core
-- | A type.
CType :: Type -> Core
-- | Introduction form for a lazily evaluated value of type Lazy T
-- for some type T. We can have multiple bindings to multiple
-- terms to create a simple target for compiling mutual recursion.
CDelay :: Bind [Name Core] [Core] -> Core
-- | Force evaluation of a lazy value.
CForce :: Core -> Core
deriving (Show, Generic, Data, Alpha)
instance Plated Core where
plate = uniplate
-- | Operators that can show up in the core language. Note that not
-- all surface language operators show up here, since some are
-- desugared into combinators of the operators here.
data Op
= -- | Addition (@+@)
OAdd
| -- | Arithmetic negation (@-@)
ONeg
| -- | Integer square root (@sqrt@)
OSqrt
| -- | Floor of fractional type (@floor@)
OFloor
| -- | Ceiling of fractional type (@ceiling@)
OCeil
| -- | Absolute value (@abs@)
OAbs
| -- | Multiplication (@*@)
OMul
| -- | Division (@/@)
ODiv
| -- | Exponentiation (@^@)
OExp
| -- | Modulo (@mod@)
OMod
| -- | Divisibility test (@|@)
ODivides
| -- | Multinomial coefficient (@choose@)
OMultinom
| -- | Factorial (@!@)
OFact
| -- | Equality test (@==@)
OEq
| -- | Less than (@<@)
OLt
| -- Type operators
-- | Enumerate the values of a type.
OEnum
| -- | Count the values of a type.
OCount
| -- Container operations
-- | Power set/bag of a given set/bag
-- (@power@).
OPower
| -- | Set/bag element test.
OBagElem
| -- | List element test.
OListElem
| -- | Map a function over a bag. Carries the
-- output type of the function.
OEachBag
| -- | Map a function over a set. Carries the
-- output type of the function.
OEachSet
| -- | Filter a bag.
OFilterBag
| -- | Merge two bags/sets.
OMerge
| -- | Bag join, i.e. union a bag of bags.
OBagUnions
| -- | Adjacency List of given graph
OSummary
| -- | Empty graph
OEmptyGraph
| -- | Construct a vertex with given value
OVertex
| -- | Graph overlay
OOverlay
| -- | Graph connect
OConnect
| -- | Map insert
OInsert
| -- | Map lookup
OLookup
| -- Ellipses
-- | Continue until end, @[x, y, z .. e]@
OUntil
| -- Container conversion
-- | set -> list conversion (sorted order).
OSetToList
| -- | bag -> set conversion (forget duplicates).
OBagToSet
| -- | bag -> list conversion (sorted order).
OBagToList
| -- | list -> set conversion (forget order, duplicates).
OListToSet
| -- | list -> bag conversion (forget order).
OListToBag
| -- | bag -> set of counts
OBagToCounts
| -- | set of counts -> bag
OCountsToBag
| -- | unsafe set of counts -> bag, assumes all are distinct
OUnsafeCountsToBag
| -- | Map k v -> Set (k × v)
OMapToSet
| -- | Set (k × v) -> Map k v
OSetToMap
| -- Number theory primitives
-- | Primality test
OIsPrime
| -- | Factorization
OFactor
| -- | Turn a rational into a (num, denom) pair
OFrac
| -- Propositions
-- | Universal quantification. Applied to a closure
-- @t1, ..., tn -> Prop@ it yields a @Prop@.
OForall [Type]
| -- | Existential quantification. Applied to a closure
-- @t1, ..., tn -> Prop@ it yields a @Prop@.
OExists [Type]
| -- | Convert Prop -> Bool via exhaustive search.
OHolds
| -- | Flip success and failure for a prop.
ONotProp
| -- | Comparison assertion
OShould BOp Type
| -- | Error for non-exhaustive pattern match
OMatchErr
| -- | Crash with a user-supplied message
OCrash
| -- | No-op/identity function
OId
| -- | Lookup OEIS sequence
OLookupSeq
| -- | Extend a List via OEIS
OExtendSeq
| -- | Not the Boolean `And`, but instead a propositional BOp
-- | Should only be seen and used with Props.
OAnd
| -- | Not the Boolean `Or`, but instead a propositional BOp
-- | Should only be seen and used with Props.
OOr
| -- | Not the Boolean `Impl`, but instead a propositional BOp
-- | Should only be seen and used with Props.
OImpl
| OSeed
| ORandom
deriving (Show, Generic, Data, Alpha, Eq, Ord)
-- | Get the arity (desired number of arguments) of a function
-- constant. A few constants have arity 0; everything else is
-- uncurried and hence has arity 1.
opArity :: Op -> Int
opArity OEmptyGraph = 0
opArity OMatchErr = 0
opArity _ = 1
substQC :: QName Core -> Core -> Core -> Core
substQC x s = transform $ \case
CVar y
| x == y -> s
| otherwise -> CVar y
t -> t
substsQC :: [(QName Core, Core)] -> Core -> Core
substsQC xs = transform $ \case
CVar y -> case P.lookup y xs of
Just c -> c
_ -> CVar y
t -> t
instance Pretty Core where
pretty = \case
CVar qn -> pretty qn
CNum r
| denominator r == 1 -> text (show (numerator r))
| otherwise -> text (show (numerator r)) <> "/" <> text (show (denominator r))
CApp (CConst op) (CPair c1 c2)
| isInfix op -> parens (pretty c1 <+> text (opToStr op) <+> pretty c2)
CApp (CConst op) c
| isPrefix op -> text (opToStr op) <> pretty c
| isPostfix op -> pretty c <> text (opToStr op)
CConst op -> pretty op
CInj s c -> withPA funPA $ selectSide s "left" "right" <+> rt (pretty c)
CCase c l r -> do
lunbind l $ \(x, lc) -> do
lunbind r $ \(y, rc) -> do
nest 2 $
"case"
<+> pretty c
<+> "of {"
$+$ vcat
[ withPA funPA $ "left" <+> rt (pretty x) <+> "->" <+> pretty lc
, withPA funPA $ "right" <+> rt (pretty y) <+> "->" <+> pretty rc
]
$+$ "}"
CUnit -> "unit"
CPair c1 c2 -> setPA initPA $ parens (pretty c1 <> ", " <> pretty c2)
CProj s c -> withPA funPA $ selectSide s "fst" "snd" <+> rt (pretty c)
CAbs _ lam -> withPA initPA $ do
lunbind lam $ \(xs, body) -> "λ" <> intercalate "," (map pretty xs) <> "." <+> lt (pretty body)
CApp c1 c2 -> withPA funPA $ lt (pretty c1) <+> rt (pretty c2)
CTest xs c -> "test" <+> prettyTestVars xs <+> pretty c
CType ty -> pretty ty
CDelay d -> withPA initPA $ do
lunbind d $ \(xs, bodies) ->
"delay" <+> intercalate "," (map pretty xs) <> "." <+> pretty (toTuple bodies)
CForce c -> withPA funPA $ "force" <+> rt (pretty c)
toTuple :: [Core] -> Core
toTuple = foldr CPair CUnit
prettyTestVars :: Members '[Reader PA, LFresh] r => [(String, Type, Name Core)] -> Sem r (Doc ann)
prettyTestVars = brackets . intercalate "," . map prettyTestVar
where
prettyTestVar (s, ty, n) = parens (intercalate "," [text s, pretty ty, pretty n])
isInfix, isPrefix, isPostfix :: Op -> Bool
isInfix OShould {} = True
isInfix op =
op
`S.member` S.fromList
[OAdd, OMul, ODiv, OExp, OMod, ODivides, OMultinom, OEq, OLt, OAnd, OOr, OImpl]
isPrefix ONeg = True
isPrefix _ = False
isPostfix OFact = True
isPostfix _ = False
instance Pretty Op where
pretty (OForall tys) = "∀" <> intercalate "," (map pretty tys) <> "."
pretty (OExists tys) = "∃" <> intercalate "," (map pretty tys) <> "."
pretty op
| isInfix op = "~" <> text (opToStr op) <> "~"
| isPrefix op = text (opToStr op) <> "~"
| isPostfix op = "~" <> text (opToStr op)
| otherwise = text (opToStr op)
opToStr :: Op -> String
opToStr = \case
OAdd -> "+"
ONeg -> "-"
OSqrt -> "sqrt"
OFloor -> "floor"
OCeil -> "ceil"
OAbs -> "abs"
OMul -> "*"
ODiv -> "/"
OExp -> "^"
OMod -> "mod"
ODivides -> "divides"
OMultinom -> "choose"
OFact -> "!"
OEq -> "=="
OLt -> "<"
OEnum -> "enumerate"
OCount -> "count"
OPower -> "power"
OBagElem -> "elem_bag"
OListElem -> "elem_list"
OEachBag -> "each_bag"
OEachSet -> "each_set"
OFilterBag -> "filter_bag"
OMerge -> "merge"
OBagUnions -> "unions_bag"
OSummary -> "summary"
OEmptyGraph -> "emptyGraph"
OVertex -> "vertex"
OOverlay -> "overlay"
OConnect -> "connect"
OInsert -> "insert"
OLookup -> "lookup"
OUntil -> "until"
OSetToList -> "set2list"
OBagToSet -> "bag2set"
OBagToList -> "bag2list"
OListToSet -> "list2set"
OListToBag -> "list2bag"
OBagToCounts -> "bag2counts"
OCountsToBag -> "counts2bag"
OUnsafeCountsToBag -> "ucounts2bag"
OMapToSet -> "map2set"
OSetToMap -> "set2map"
OIsPrime -> "isPrime"
OFactor -> "factor"
OFrac -> "frac"
OHolds -> "holds"
ONotProp -> "not"
OShould Eq _ -> "=!="
OShould Neq _ -> "=!!="
OShould Lt _ -> "!<"
OShould Gt _ -> "!>"
OShould Leq _ -> "!<="
OShould Geq _ -> "!>="
OShould Divides _ -> "!|"
OShould _ _ -> "<!>"
OMatchErr -> "matchErr"
OCrash -> "crash"
OId -> "id"
OLookupSeq -> "lookupSeq"
OExtendSeq -> "extendSeq"
OForall {} -> "∀"
OExists {} -> "∃"
OAnd -> "and"
OOr -> "or"
OImpl -> "implies"
ORandom -> "random"
OSeed -> "seed"