Run newtype optimisation as part of lower (#210)

- Introduce the concept of "type representations". These are a concrete
   part of the (core) module signature, rather than provided by an
   optimisation.

 - Lower is now guided by these type representations. We will emit
   method calls instead of constructor invocations for type wrappers,
   and the pattern compiler will also generate the appropriate code.

Closes #196

amuletml.cabal

@@ -281,6 +281,7 @@ library
                      , Core.Lower
                      , Core.Lower.Basic
                      , Core.Lower.Pattern
+                     , Core.Lower.TypeRepr
                      , Core.Types
                      , Core.Builtin
                      , Core.Optimise
@@ -291,7 +292,6 @@ library
                      , Core.Optimise.Reduce.Inline
                      , Core.Optimise.Reduce.Pattern
                      , Core.Optimise.Sinking
-                     , Core.Optimise.Newtype
                      , Core.Optimise.Uncurry
                      , Core.Optimise.DeadCode
                      , Core.Optimise.CommonExpElim

bin/Amc.hs

@@ -22,7 +22,6 @@ import Backend.Lua
 import qualified Syntax.Builtin as Bi
 
 import Core.Optimise.Reduce (reducePass)
-import Core.Optimise.Newtype (killNewtypePass)
 import Core.Optimise.DeadCode (deadCodePass)
 import Core.Simplify (optimise)
 import Core.Core (Stmt)
@@ -65,7 +64,7 @@ runCompile opt (DoLint lint) dconfig file = do
             Opt -> optimise lint core
             NoOpt -> do
               lintIt "Lower" (checkStmt emptyScope core) (pure ())
-              (lintIt "Optimised"  =<< checkStmt emptyScope) . deadCodePass <$> (reducePass =<< killNewtypePass core)
+              (lintIt "Optimised"  =<< checkStmt emptyScope) . deadCodePass <$> reducePass core
           lua = compileProgram optimised
       in ( Just (env, core, optimised, lua)
          , errors

src/Core/Lower.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE LambdaCase, TupleSections,
+{-# LANGUAGE TupleSections,
     PatternSynonyms, RankNTypes, ScopedTypeVariables, FlexibleContexts,
     ConstraintKinds, OverloadedStrings, TypeFamilies #-}
 module Core.Lower
@@ -32,6 +32,7 @@ import Core.Optimise (substituteInType, substituteInTys, fresh, freshFrom)
 import Core.Core hiding (Atom, Term, Stmt, Type, Pattern, Arm)
 import Core.Core (pattern Atom)
 import Core.Types (unify, unifyClosed, replaceTy)
+import Core.Lower.TypeRepr
 import Core.Lower.Pattern
 import Core.Lower.Basic
 import Core.Var
@@ -40,7 +41,7 @@ import qualified Syntax as S
 import Syntax.Let
 import Syntax.Var (Var, Typed, VarResolved(..))
 import Syntax.Transform
-import Syntax (Expr(..), Pattern(..), Skolem(..), ModuleTerm(..), Toplevel(..), Constructor(..), Arm(..))
+import Syntax (Expr(..), Pattern(..), Skolem(..), ModuleTerm(..), Toplevel(..), Arm(..))
 
 import Text.Pretty.Semantic (pretty)
 
@@ -52,13 +53,18 @@ type Stmt = C.Stmt CoVar
 type Lower = ContT Term
 
 defaultState :: LowerState
-defaultState = LS mempty ctors  mempty where
+defaultState = LS mempty ctors types where
   ctors :: VarMap.Map (C.Type CoVar)
   ctors = VarMap.fromList
            [ (C.vCONS,
               ForallTy (Relevant name) StarTy $
                 VarTy name `prodTy` AppTy C.tyList (VarTy name) `arrTy` AppTy C.tyList (VarTy name))
            , (C.vNIL, ForallTy (Relevant name) StarTy $ AppTy C.tyList (VarTy name))]
+  types :: VarMap.Map TypeRepr
+  types = VarMap.fromList
+            ( (C.vList, SumTy (VarSet.fromList [C.vCONS, C.vNIL]))
+            : map (,OpaqueTy) [ C.vBool, C.vInt, C.vString, C.vFloat, C.vUnit
+                              , C.vLazy, C.vArrow, C.vProduct, C.vRefTy ] )
   name = C.tyvarA
   arrTy = ForallTy Irrelevant
   prodTy a b = RowsTy NilTy [("_1", a), ("_2", b)]
@@ -365,22 +371,11 @@ lowerProg' (LetStmt _ vs:prg) = do
     vs' <- lowerLet vs
     foldr ((.) . ((:) . C.StmtLet)) id vs' <$$> lowerProg' prg
 
-lowerProg' (TypeDecl _ var _ Nothing _:prg) =
-  (C.Type (mkType var) []:) <$$> lowerProg' prg
-lowerProg' (TypeDecl _ var _ (Just cons) _:prg) = do
-  let cons' = map (\case
-                       UnitCon _ p (_, t) -> (p, mkCon p, lowerType t)
-                       ArgCon _ p _ (_, t) -> (p, mkCon p, lowerType t)
-                       GadtCon _ p t _ -> (p, mkCon p, lowerType t))
-                cons
-      ccons = map (\(_, a, b) -> (a, b)) cons'
-      scons = map (\(a, _, b) -> (mkCon a, b)) cons'
-
-      conset = VarSet.fromList (map fst scons)
-
-  (C.Type (mkType var) ccons:) <$$> local (\s ->
-    s { ctors = VarMap.union (VarMap.fromList scons) (ctors s)
-      , types = VarMap.insert (mkType var) conset (types s)
+lowerProg' (TypeDecl _ var _ cons  _:prg) = do
+  ~(tyStmts@(C.Type _ cs:_), repr) <- getTypeRepr (mkType var) cons
+  (tyStmts++) <$$> local (\s ->
+    s { ctors = VarMap.union (VarMap.fromList cs) (ctors s)
+      , types = VarMap.insert (mkType var) repr (types s)
       }) (lowerProg' prg)
 
 lowerLet :: MonadLower m => [S.Binding Typed] -> m [Binding CoVar]

src/Core/Lower/Basic.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE OverloadedStrings, ConstraintKinds, FlexibleContexts #-}
 module Core.Lower.Basic
-  ( LowerState(..)
+  ( TypeRepr(..)
+  , LowerState(..)
   , LowerTrack
   , MonadLower
   , mkTyvar, mkVal, mkType, mkCo, mkCon, mkVar
@@ -24,12 +25,23 @@ import qualified Syntax as S
 import Syntax.Var (VarResolved(..), Var, Resolved, Typed)
 import Syntax (Lit(..), Skolem(..))
 
+data TypeRepr
+  = OpaqueTy -- ^ An opaque type, for interfacing with foreign values.
+  | SumTy VarSet.Set -- ^ A sum type, with the set of constructors.
+
+  -- | A type which just wraps another.
+  --
+  -- This holds the name of the constructor, and the inner and outer
+  -- type, both sharing their free variables.
+  | WrapperTy CoVar (C.Type CoVar) (C.Type CoVar)
+  deriving (Show, Eq)
+
 data LowerState
   = LS
     { vars  :: VarMap.Map (C.Type CoVar)
     , ctors :: VarMap.Map (C.Type CoVar)
     -- | The map of types to their constructors /if they have any/.
-    , types :: VarMap.Map VarSet.Set
+    , types :: VarMap.Map TypeRepr
     } deriving (Eq, Show)
 
 instance Semigroup LowerState where

src/Core/Lower/Pattern.hs

@@ -31,6 +31,7 @@ import Data.Bifunctor
 import Data.Foldable
 import Data.Triple
 import Data.Maybe
+import Data.Span
 
 import qualified Core.Core as C
 import Core.Optimise (substituteInType, substituteInTys, fresh, freshFrom)
@@ -85,13 +86,23 @@ data ArmNode
     , nodeNodes   :: [(Pattern CoVar, ArmNode)]
       -- ^ The child nodes, and their associated pattern.
     }
+  | ArmLet
+    { nodeArms    :: ArmSet
+    , nodeSuccess :: [ArmLeaf] -- Should be empty
+    , nodeBind    :: (CoVar, Type CoVar, Term CoVar)
+    , nodeBody    :: ArmNode
+    }
   deriving (Show)
 
 instance Pretty ArmNode where
   pretty (ArmComplete arms success) = "Complete" <+> parens (shown arms) <+> shown success
   pretty (ArmMatch arms success atom nodes)
     = "Match" <+> parens (shown arms) <+> shown success <+> pretty atom
     <#> (indent 2 . vsep $ map (\(p, n) -> pretty p <+> "=>" <#> indent 2 (pretty n)) nodes)
+  pretty (ArmLet arms success (v, ty, x) node)
+    = "Let" <+> parens (shown arms) <+> shown success <+> pretty v <+> colon <+> pretty ty <+> equals <+> pretty x
+    <#> indent 2 (pretty node)
+
 
 -- | A of a single case in a match expression.
 data PatternRow
@@ -197,6 +208,8 @@ flattenNode bodies guards (ArmMatch _ leafs atom' children) = do
       let branches = foldr (flip (HSet.foldr add) . nodeArms . snd) mempty children
       in foldr (add . leafArm) branches leafs
       where add k = HMap.insertWith (+) k (1 :: Int)
+flattenNode bodies guards (ArmLet _ _ bind child) =
+  C.Let (One bind) <$> flattenNode bodies guards child
 
 -- | Lift a pattern match into a lambda, passing arguments as values.
 generateBinds :: forall m. MonadLower m
@@ -359,12 +372,11 @@ lowerOne tys rss = do
       getCtors v = do
         ctor <- VarMap.lookup v (ctors state)
         ty <- getType ctor
-        VarMap.lookup ty (types state)
-
-      getType (ForallTy _ _ t) = getType t
-      getType (ConTy a) = pure a
-      getType (AppTy f _) = getType f
-      getType _ = Nothing
+        case VarMap.lookup ty (types state) of
+          Nothing -> error ("Cannot find " ++ show ty)
+          Just OpaqueTy -> Nothing
+          Just (SumTy ctors) -> Just ctors
+          Just (WrapperTy ctor _ _) -> Just (VarSet.singleton ctor)
 
     -- | Compute the "arity" heuristic for a given row variable.
     --
@@ -415,8 +427,8 @@ lowerOneOf preLeafs var ty tys = go [] . map prepare
 
     go unc [] = lowerOne tys (reverse unc)
     go unc rs@((S.PRecord{},_):_) = goRows unc mempty rs
-    go unc rs@((S.Destructure{},_):_) = goCtors unc mempty rs
-    go unc rs@((S.PGadtCon{},_):_) = goCtors unc mempty rs
+    go unc rs@((S.Destructure{},_):_) = goCtorsWith unc rs
+    go unc rs@((S.PGadtCon{},_):_) = goCtorsWith unc rs
     go unc rs@((S.PLiteral{},_):_) = goLiterals unc mempty rs
     go unc ((p, r):rs) = go (goGeneric p r:unc) rs
 
@@ -464,6 +476,46 @@ lowerOneOf preLeafs var ty tys = go [] . map prepare
         pure ( Map.insert f (v, lowerType (S.getType p)) fs
              , VarMap.insert v p ps )
 
+    goCtorsWith unc rs = do
+     let Just tyName = getType ty
+     repr <- asks (fromMaybe (error ("Cannot find " ++ show tyName)) . VarMap.lookup tyName . types)
+     case repr of
+       OpaqueTy -> error "Impossible matching on opaque type"
+       SumTy _ -> goCtors unc mempty rs
+       WrapperTy _ from to -> do
+         let Just map = unify to ty
+             from' = substituteInType map from
+         coVar <- case rs of
+               ((S.PGadtCon _ _ _ (Just child) _, _):_) -> freshFromPat child
+               _ -> fresh ValueVar
+         node <- goNewtype unc (Capture coVar from') rs
+         pure (ArmLet (nodeArms node) mempty
+                (coVar, from', Cast (Ref var ty) from' (SameRepr ty from'))
+                node)
+
+    -- | Split patterns into those matching against the constructor and those not
+    goNewtype :: [PatternRow] -> Capture CoVar
+              -> [(S.Pattern Typed, PatternRow)]
+              -> m ArmNode
+    goNewtype unc (Capture c cty) [] =
+      lowerOne (VarMap.insert c cty tys) (reverse unc)
+
+    goNewtype unc cap@(Capture c _) (( S.PGadtCon _ _ [] (Just p) _
+                                     , PR arm pats gd vBind tyBind ):rs) =
+      -- The wrapped value is matched by the pattern - focus on that next.
+      let r' = PR arm (VarMap.insert c p pats) gd vBind tyBind
+      in goNewtype (r':unc) cap rs
+
+    goNewtype unc cap@(Capture c _) (( S.PGadtCon _ _ [(v, t)] Nothing _
+                                     , PR arm pats gd vBind tyBind ):rs) =
+      -- The wrapped value is the dictionary - just add a wildcard pattern.
+      let r' = PR arm (VarMap.insert c (S.Capture v (internal, t)) pats) gd vBind tyBind
+      in goNewtype (r':unc) cap rs
+
+    goNewtype _ _ ((S.PGadtCon{}, _):_) = error "Impossible: Malformed pattern for newtype."
+
+    goNewtype unc cap ((p, r):rs) = goNewtype (goGeneric p r:unc) cap rs
+
     -- | Build up a mapping of (constructors -> (contents variable, rows)).
     goCtors :: [PatternRow] -> VarMap.Map ([Capture CoVar], [PatternRow])
             -> [(S.Pattern Typed, PatternRow)]
@@ -577,3 +629,9 @@ dropNForalls :: Int -> Type a -> Type a
 dropNForalls 0 t = t
 dropNForalls x (ForallTy _ _ t) = dropNForalls (x - 1) t
 dropNForalls _ _ = undefined
+
+getType :: Type a -> Maybe a
+getType (ForallTy _ _ t) = getType t
+getType (ConTy a) = pure a
+getType (AppTy f _) = getType f
+getType _ = Nothing

src/Core/Lower/TypeRepr.hs

@@ -0,0 +1,62 @@
+{-# LANGUAGE LambdaCase, ScopedTypeVariables #-}
+module Core.Lower.TypeRepr (getTypeRepr) where
+
+import qualified Data.VarSet as VarSet
+
+import Control.Monad.Namey
+import Control.Monad
+
+import Core.Lower.Basic
+import Core.Optimise
+
+import Syntax (Constructor(..))
+import Syntax.Var (Typed)
+
+getTypeRepr :: MonadNamey m
+            => CoVar -> Maybe [Constructor Typed]
+            -> m ([Stmt CoVar], TypeRepr)
+getTypeRepr var Nothing = pure ([Type var []], OpaqueTy)
+getTypeRepr var (Just ctors) =
+  let ctors' = map (\case
+                       UnitCon _ p (_, t)  -> (mkCon p, lowerType t)
+                       ArgCon _ p _ (_, t) -> (mkCon p, lowerType t)
+                       GadtCon _ p t _     -> (mkCon p, lowerType t)) ctors
+  in case ctors' of
+    [(ctor, ty)] | Just nt@(Spine _ dom cod) <- isNewtype ty -> do
+      let CoVar name id _ = ctor
+
+      wrapper <- newtypeWorker nt
+      pure ( [ Type var [], StmtLet (One (CoVar name id ValueVar, ty, wrapper))]
+           , WrapperTy ctor dom cod )
+
+    _ -> pure ( [ Type var ctors' ]
+              , SumTy (VarSet.fromList (map fst ctors')) )
+
+isNewtype :: IsVar a => Type a -> Maybe (Spine a)
+isNewtype (ForallTy Irrelevant _ ForallTy{}) = Nothing -- Cannot have multiple relevant arguments
+isNewtype (ForallTy Irrelevant from to) =
+  pure (Spine [(Irrelevant, from)] from to)
+isNewtype (ForallTy (Relevant var) k rest) = do
+  (Spine tys from to) <- isNewtype rest
+  guard (var `occursInTy` to)
+  pure (Spine ((Relevant var, k):tys) from to)
+isNewtype _ = Nothing
+
+data Spine a =
+  Spine [(BoundTv a, Type a)] (Type a) (Type a)
+  deriving (Eq, Show, Ord)
+
+newtypeWorker :: forall a m. (IsVar a, MonadNamey m)
+              => Spine a -> m (Term a)
+newtypeWorker (Spine tys dom cod) = do
+  let wrap :: [(BoundTv a, Type a)] -> (a -> Type a -> Term a) -> m (Term a)
+      wrap ((Relevant v, c):ts) ex = Lam (TypeArgument v c) <$> wrap ts ex
+      wrap [(Irrelevant, c)] ex = do
+        v <- fresh ValueVar
+        Lam (TermArgument (fromVar v) c) <$> pure (ex (fromVar v) c)
+      wrap _ _ = undefined
+
+      work :: a -> Type a -> Term a
+      work var ty = Cast (Ref var ty) cod (SameRepr dom cod)
+
+  wrap tys work