Enumerate.hs

{-# LANGUAGE NondecreasingIndentation #-}

-- |
-- Module      :  Disco.Enumerate
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Enumerate values inhabiting Disco types.
module Disco.Enumerate (
  ValueEnumeration,

  -- * Base types
  enumVoid,
  enumUnit,
  enumBool,
  enumN,
  enumZ,
  enumF,
  enumQ,
  enumC,

  -- * Containers
  enumSet,
  --  , enumBag
  enumList,

  -- * Any type
  enumType,
  enumTypes,

  -- * Lifted functions that return lists
  enumerateType,
  enumerateTypes,
)
where

import qualified Data.Enumeration.Invertible as E
import Disco.AST.Generic (Side (..))
import Disco.Types
import Disco.Value

type ValueEnumeration = E.IEnumeration Value

-- | Enumerate all values of type @Void@ (none).
enumVoid :: ValueEnumeration
enumVoid = E.void

-- | Enumerate all values of type @Unit@ (the single value @unit@).
enumUnit :: ValueEnumeration
enumUnit = E.singleton VUnit

-- | Enumerate the values of type @Bool@ as @[false, true]@.
enumBool :: ValueEnumeration
enumBool = E.mapE toV fromV $ E.finiteList [L, R]
 where
  toV i = VInj i VUnit
  fromV (VInj i VUnit) = i
  fromV _ = error "enumBool.fromV: value isn't a bool"

-- | Enumerate all values of type @Nat@ (0, 1, 2, ...).
enumN :: ValueEnumeration
enumN = E.mapE (ratv . fromInteger) (floor . vrat) E.nat

-- | Enumerate all values of type @Integer@ (0, 1, -1, 2, -2, ...).
enumZ :: ValueEnumeration
enumZ = E.mapE (ratv . fromInteger) (floor . vrat) E.int

-- | Enumerate all values of type @Fractional@ in the Calkin-Wilf
--   order (1, 1/2, 2, 1/3, 3/2, 2/3, 3, ...).
enumF :: ValueEnumeration
enumF = E.mapE ratv vrat E.cw

-- | Enumerate all values of type @Rational@ in the Calkin-Wilf order,
--   with negatives interleaved (0, 1, -1, 1/2, -1/2, 2, -2, ...).
enumQ :: ValueEnumeration
enumQ = E.mapE ratv vrat E.rat

-- | Enumerate all Unicode characters.
enumC :: ValueEnumeration
enumC = E.mapE toV fromV (E.boundedEnum @Char)
 where
  toV = ratv . fromIntegral . fromEnum
  fromV = toEnum . floor . vrat

-- | Enumerate all *finite* sets over a certain element type, given an
--   enumeration of the elements.  If we think of each finite set as a
--   binary string indicating which elements in the enumeration are
--   members, the sets are enumerated in order of the binary strings.
enumSet :: ValueEnumeration -> ValueEnumeration
enumSet e = E.mapE toV fromV (E.finiteSubsetOf e)
 where
  toV = VBag . map (,1)
  fromV (VBag vs) = map fst vs
  fromV _ = error "enumSet.fromV: value isn't a set"

-- | Enumerate all *finite* lists over a certain element type, given
--   an enumeration of the elements.  It is very difficult to describe
--   the order in which the lists are generated.
enumList :: ValueEnumeration -> ValueEnumeration
enumList e = E.mapE toV fromV (E.listOf e)
 where
  toV = foldr VCons VNil
  fromV (VCons h t) = h : fromV t
  fromV VNil = []
  fromV _ = error "enumList.fromV: value isn't a list"

-- | Enumerate all functions from a finite domain, given enumerations
--   for the domain and codomain.
enumFunction :: ValueEnumeration -> ValueEnumeration -> ValueEnumeration
enumFunction xs ys =
  case (E.card xs, E.card ys) of
    (E.Finite 0, _) -> E.singleton (VFun $ \_ -> error "enumFunction: void function called")
    (_, E.Finite 0) -> E.void
    (_, E.Finite 1) -> E.singleton (VFun $ \_ -> E.select ys 0)
    _ -> E.mapE toV fromV (E.functionOf xs ys)
 where
  -- XXX TODO: better error message on functions with an infinite domain

  toV = VFun
  fromV (VFun f) = f
  fromV _ = error "enumFunction.fromV: value isn't a VFun"

-- | Enumerate all values of a product type, given enumerations of the
--   two component types.  Uses a fair interleaving for infinite
--   component types.
enumProd :: ValueEnumeration -> ValueEnumeration -> ValueEnumeration
enumProd xs ys = E.mapE toV fromV $ (E.><) xs ys
 where
  toV (x, y) = VPair x y
  fromV (VPair x y) = (x, y)
  fromV _ = error "enumProd.fromV: value isn't a pair"

-- | Enumerate all values of a sum type, given enumerations of the two
--   component types.
enumSum :: ValueEnumeration -> ValueEnumeration -> ValueEnumeration
enumSum xs ys = E.mapE toV fromV $ (E.<+>) xs ys
 where
  toV (Left x) = VInj L x
  toV (Right y) = VInj R y
  fromV (VInj L x) = Left x
  fromV (VInj R y) = Right y
  fromV _ = error "enumSum.fromV: value isn't a sum"

-- | Enumerate the values of a given type.
enumType :: Type -> ValueEnumeration
enumType TyVoid = enumVoid
enumType TyUnit = enumUnit
enumType TyBool = enumBool
enumType TyN = enumN
enumType TyZ = enumZ
enumType TyF = enumF
enumType TyQ = enumQ
enumType TyC = enumC
enumType (TySet t) = enumSet (enumType t)
enumType (TyList t) = enumList (enumType t)
enumType (a :*: b) = enumProd (enumType a) (enumType b)
enumType (a :+: b) = enumSum (enumType a) (enumType b)
enumType (a :->: b) = enumFunction (enumType a) (enumType b)
enumType ty = error $ "enumType: can't enumerate " ++ show ty

-- | Enumerate a finite product of types.
enumTypes :: [Type] -> E.IEnumeration [Value]
enumTypes [] = E.singleton []
enumTypes (t : ts) = E.mapE toL fromL $ (E.><) (enumType t) (enumTypes ts)
 where
  toL (x, xs) = x : xs
  fromL (x : xs) = (x, xs)
  fromL [] = error "enumTypes.fromL: empty list not in enumeration range"

-- | Produce an actual list of the values of a type.
enumerateType :: Type -> [Value]
enumerateType = E.enumerate . enumType

-- | Produce an actual list of values enumerated from a finite product
--   of types.
enumerateTypes :: [Type] -> [[Value]]
enumerateTypes = E.enumerate . enumTypes