ToneInference.hs

{-# LANGUAGE FlexibleInstances #-}
module MiniToneInference where

import Prelude hiding ((&&))
import Control.Applicative
import Control.Monad
import Data.Map.Strict (Map, (!))
import Data.Maybe (fromJust)
import Data.Monoid ((<>))
import Data.Set (isSubsetOf)
import qualified Data.Bool as Bool
import qualified Data.Map.Strict as M

-- Laws: (<:) is reflexive and transitive.
class Preorder a where (<:) :: a -> a -> Bool

-- Laws:
-- 1. (a <: b) iff (a && b == a).
-- 2. (&&) is idempotent, commutative, associative, and has top as identity.
class Preorder a => MeetSemilattice a where
  top :: a
  (&&) :: a -> a -> a
  meet :: [a] -> a
  meet = foldr (&&) top

infix  4 <:
infixr 3 &&
instance Preorder Bool where False <: True = False; x <: y = True
instance MeetSemilattice Bool where (&&) = (Bool.&&); top = True; meet = and


---------- Tones ----------
-- Tones modify the ordering on a preorder A:
--
-- 1. Id A = A.
--
-- 2. (x <= y : Op A) iff (y <= x : A).
--
-- 3. (x <= y : Iso A) iff (x <= y and y <= x : A).
--
-- 4. (x <= y : Path A) if (x <= y or y <= x : A). This is an "if", not an "if
-- and only if"; in general, (x <= y or y <= x) isn't transitive, so we have to
-- take its transitive closure. Because of this, actually *testing* inequality
-- at (Path A) is difficult for this reason. Luckily, we won't need to do that.
--
-- I use variables t,u,v to stand for tones. Tones have the following property:
-- if (f : A -> B) is monotone, then so is (f : t A -> t B) for any tone t.
-- This makes them endofunctors on the category of preorders & monotone maps.
data Tone = Id | Op | Iso | Path deriving (Show, Eq, Ord)

-- For preorders A, B, let (A <: B) iff A is a subpreorder of B; that is, if A ⊆
-- B and x ≤ y : A implies x ≤ y : B. Then tones form a lattice as follows: let
-- (t <= u) iff (t A <: u A) for all preorders A. The lattice is diamond-shaped:
--
--          Path                greatest: ∀A. (A <: Path)
--         /   \
--        Id   Op
--         \   /
--          Iso                 least: ∀A. (Iso <: A)
--
instance Preorder Tone where
  Iso <: _ = True
  _ <: Path = True
  x <: y | (x == y) = True
  _ <: _ = False

instance MeetSemilattice Tone where
  top = Path
  x && y | x <: y = x
         | y <: x = y
         | otherwise = Iso

-- Tones, as endofunctors on Preorder, form a monoid under composition.
-- I give composition in the same order as functions: for a preorder A,
-- (t <> u) A == t (u A).
instance Monoid Tone where
  mempty = Id
  mappend x Id = x
  mappend Id y = y
  mappend Op Op = Id
  -- the order of the remaining clauses matters.
  mappend x Iso  = Iso   -- hi priority
  mappend x Path = Path  -- hi priority
  mappend Iso  x = Iso   -- lo priority
  mappend Path y = Iso   -- lo priority

-- (&&) and (<>) form a semiring; that is, (<>) distributes over (&&):
--
--     t <> (u && v) == (t <> u) && (t <> v)
--     (t && u) <> v == (t <> v) && (u <> v)
--
-- A number of recent and not-so-recent type systems also use semiring
-- annotations on variables, including:
--
-- 1. "Linear Haskell" (POPL 2018)
-- 2. "I Got Plenty o' Nuttin" (McBride)
-- 3. "Monotonicity Types" (PMLDC 2017)
-- 4. "Bounded Linear Types in a Resource Semiring" (ESOP 2014)
-- 5. "The Next 700 Modal Type Assignment Systems" (TYPES 2015)
-- 6. Petricek, Orchard, et al's work on coeffects
-- 7. Provenance and information flow typing, possibly?
-- 8. "A Fibrational Framework for Substructural and Modal Logics" (FSCD 2017),
-- which massively generalizes this pattern.


---------- Types ----------
infixr 4 :->
data Type
  = Bool                  -- booleans
  | Prod [Type]           -- Tuples or product types
  | Sum (Map String Type) -- Sum types; values tagged with a constructor
  | (:->) Type Type       -- Functions

  -- Finally, a type can be wrapped with a tone. This doesn't change the type's
  -- values, but it alters how they're ordered for purposes of checking
  -- monotonicity. Datafun's type system ensures all functions are monotone, so
  -- to use something "non-monotonically", we have to wrap it with a tone.
  | Mode Tone Type
  -- invariant: in (Mode t a), t /= Path.
  -- This is because I don't know how to handle the (t = Path) case yet.
  -- I'm not sure there is a "good" way to do it.
    deriving Show

-- Returns True only if a type has a least element.
hasBottom :: Tone -> Type -> Bool
hasBottom t (Mode s a) = hasBottom (t <> s) a
hasBottom Id Bool = True
hasBottom Op Bool = True
hasBottom _  Bool = False
hasBottom t (Prod as) = all (hasBottom t) as
hasBottom t (Sum h) | [(_,a)] <- M.toList h = hasBottom t a
hasBottom _ (Sum _) = False
hasBottom t (a :-> b) = hasBottom t b


---------- Subtyping ----------
-- (leftAdjoint t == u) when u is left adjoint to t, meaning that ∀ preorders
-- A, B and functions f : A -> B:
--
--     f is monotone from (A -> t B)
--                  iff
--     f is monotone from (u A -> B)
--
-- This property is useful for type-checking terms that use variables of type
-- (Mode t a), and in the subtyping rules for (Mode t a) as well.
--
-- Exercise for the reader: why does Path have no left adjoint?
leftAdjoint :: Tone -> Tone
leftAdjoint Path = error "Path has no left adjoint"
leftAdjoint Iso = Path
leftAdjoint Id = Id
leftAdjoint Op = Op

-- stripTone t a = (u,b) such that (Mode t a) and (Mode u b) are equivalent,
-- and b /= (Mode _ _).
stripTone :: Tone -> Type -> (Tone, Type)
stripTone t (Mode u b) = stripTone (t <> u) b
stripTone t a = (t,a)

-- (detune a == (t,b)) when (t a <: b) and b /= (Mode _ _). Used when coercing
-- something into a non-modal type - e.g. function application coerces the
-- applied expression to function type.
detune :: Type -> (Tone, Type)
detune a = (leftAdjoint t, b) where (t,b) = stripTone Id a

-- (subtype a b == Just t) implies (t a <: b), and t is the greatest tone for
-- which this is true.
subtype :: (Monad m, Alternative m) => Type -> Type -> m Tone
-- Mode introduction/elimination.
subtype a (Mode t b) = (t <>) <$> subtype a b
subtype (Mode t a) b = (<> leftAdjoint t) <$> subtype a b
-- More interesting types.
subtype Bool Bool = pure Id
subtype (Prod as) (Prod bs)
  | length as == length bs = meet <$> zipWithM subtype as bs
subtype (Sum as) (Sum bs)
  | M.keysSet as `isSubsetOf` M.keysSet bs =
    meet <$> sequence [subtype a (bs ! name) | (name, a) <- M.toList as]
subtype (a1 :-> b1) (a2 :-> b2) = do
  -- Exercise for the reader: justify this code, esp. the `guard` expression.
  dom <- subtype a2 a1
  cod <- subtype b1 b2
  guard $ case cod of Path -> Path == (Path <> dom)
                      _ -> leftAdjoint cod <: dom
  pure cod
-- Failure cases.
subtype Bool _ = fail "nope"
subtype Prod{} _ = fail "nope"
subtype Sum{} _ = fail "nope"
subtype (:->){} _ = fail "nope"


---------- Expressions and patterns ----------
type Var = String               -- a variable name
type Tag = String               -- a "tag", or constructor of a sum type

-- Patterns can be variables, tuples, or constructors.
data Pat = PVar Var | PTuple [Pat] | PTag Tag Pat deriving Show

-- Expressions whose types can be inferred:
data Expr = Var Var             -- variables,
          | LitBool Bool        -- literals,
          | The Type Term       -- any type-annotated term,
          | App Expr Term       -- and function applications.
  deriving Show

-- Terms, whose types can be checked:
data Term = Expr Expr              -- any expression whose type can be inferred,
          | If Term Term Term      -- conditionals,
          | When Term Term         -- monotone conditionals,
          | Lambda Var Term        -- functions,
          | Tag Tag Term           -- tagged values of a sum type,
          | Tuple [Term]           -- tuples,
          | Let Var Expr Term      -- let-bindings,
          | Case Expr [(Pat,Term)] -- and case-expressions.
  deriving Show

-- We could probably expand the set of inferrable expressions. But for the sake
-- of clarity, I don't, to avoid redundancy. For example, "let x = e1 in e2" can
-- infer if "e2" can infer; but then we need two cases for typechecking let - an
-- inferring case, and a checking case.


---------- Type checking ----------
-- Type & tone contexts. Type contexts map variables to their types; tone
-- contexts, to the tones at which they're used.
type Types = Map Var Type
type Tones = Map Var Tone

-- Composing a tone over a tone context. For example, if an expression `e` uses
-- variables at tones given by `ts`; and we use `e` with tone `t`; then overall
-- we use those variables at tones given by (t <@ ts).
(<@) :: Tone -> Tones -> Tones
t <@ ts = fmap (t <>) ts

-- Taking the meet of two tone contexts. Absent variables are regarded as if
-- mapped to Path.
instance MeetSemilattice Tones where top = M.empty; (&&) = M.unionWith (&&)
instance Preorder Tones where (<:) = error "Left as exercise for the reader."

-- (checkPat t a p == cx) means pattern `p` can match a value of type `a` at
-- tone `t`, producing variables with types in `cx`.
checkPat :: Tone -> Type -> Pat -> Types
checkPat t (Mode u b) p = checkPat (t <> u) b p
checkPat t a (PVar x) = M.singleton x (Mode t a)
checkPat t (Prod as) (PTuple ps)
  | length as == length ps = M.unionsWith uhoh $ zipWith (checkPat t) as ps
  where uhoh _ _ = error "cannot use same variable twice in pattern"
checkPat t (Sum arms) (PTag name p) = checkPat t (arms ! name) p
checkPat t _ PTuple{} = error "tuple must have product type"
checkPat t _ PTag{} = error "tagged pattern must have sum type"

-- (infer cx e == (a,ts)) means `e` infers type `a` using variables of types in
-- `cx` at tones given in `ts`.
infer :: Types -> Expr -> (Type, Tones)
infer cx (Var x) = (cx ! x, M.singleton x Id)
infer cx (LitBool b) = (Bool, M.empty)
infer cx (The a m) = (a, check cx a m)
infer cx (App e m) = let (tp, ts) = infer cx e
                         (t, a :-> b) = detune tp
                     in (b, (t <@ ts) && check cx a m)

-- (check cx a m == ts) means `e` checks at type `a` using variables of types
-- in `cx` at tones given in `ts`.
check :: Types -> Type -> Term -> Tones
-- first, deal with modes on the type.
check cx (Mode t a) m = t <@ check cx a m
-- then, examine the term.
check cx a (Expr e) = let (b,ts) = infer cx e
                      in fromJust (subtype a b) <@ ts
check cx a (If m n o) =
  check cx (Mode Iso Bool) m && check cx a n && check cx a o
check cx a (When m n)
  | hasBottom Id a = check cx Bool m && check cx a m
  | otherwise = error "use of `when` at non-pointed type"
check cx (a :-> b) (Lambda x m) =
  let ts = check (M.insert x a cx) b m
  in if Id <: ts ! x then M.delete x ts
     else error ("non-monotone use of variable " ++ x)
check cx (Sum as) (Tag name m) = check cx (as ! name) m
check cx (Prod as) (Tuple ms)
  | length as == length ms = meet $ zipWith (check cx) as ms
check cx expected (Let x e m) =
  let (a, eTones) = infer cx e
      mTones = check (M.insert x a cx) expected m
      xTone = M.findWithDefault Path x mTones
  in M.delete x mTones && (xTone <@ eTones)
check cx expected (Case e arms) =
  let (subjType, subjTones) = infer cx e in
  meet $ flip map arms $ \(p, m) ->
    let patCx = checkPat Id subjType p
        -- The monoid instance on Map is a left-biased union.
        armTones = check (patCx <> cx) expected m
        -- Split tones: are they on variables from subject, or external?
        (patTones, cxTones) =
          M.partitionWithKey (\x _ -> M.member x patCx) armTones
        -- The tone at which we use the subject.
        useTone = meet $ map snd $ M.toList patTones
    in cxTones && (useTone <@ subjTones)
-- Failure cases
check cx _ Lambda{} = error "lambda must have function type"
check cx _ Tag{} = error "tagged expression must have sum type"
check cx _ Tuple{} = error "tuple must have product type"