@@ -5,7 +5,8 @@ Description : Coordinates the compilation process.
55License : BSD3
66Maintainer : The Idris Community.
77-}
8−{-# LANGUAGE CPP, FlexibleContexts, PatternGuards, TypeSynonymInstances #-}
8+{-# LANGUAGE CPP, FlexibleContexts, MultiWayIf, NamedFieldPuns, PatternGuards,
9+ TypeSynonymInstances #-}
910
1011module IRTS.Compiler(compile, generate) where
1112
@@ -34,6 +35,7 @@ import Control.Category
3435import Control.Monad.State
3536import Data.List
3637import qualified Data.Map as M
38+import Data.Maybe (maybe)
3739import Data.Ord
3840import qualified Data.Set as S
3941import System.Directory
@@ -393,6 +395,16 @@ irTerm top vs env tm@(App _ f a) = do
393395 EQ | isNewtype
394396 -> irTerm top vs env (head argsPruned)
395397
398+ -- compile Nat-likes as bigints
399+ | Just LikeZ <- isLikeNat ist n
400+ -> irTerm top vs env $ Constant (BI 0)
401+
402+ -- compile Nat-likes as bigints
403+ | Just LikeS <- isLikeNat ist n
404+ -> irTerm top vs env $ mkApp
405+ (P Ref (sUN "prim__addBigInt") Erased)
406+ (Constant (BI 1) : argsPruned)
407+
396408 | otherwise -- not newtype, plain data ctor
397409 -> buildApp (LV n) argsPruned
398410
@@ -405,6 +417,14 @@ irTerm top vs env tm@(App _ f a) = do
405417 | isNewtype -- newtype but the value is not among args yet
406418 -> return . padLams $ \[vn] -> LApp False (LV n) [LV vn]
407419
420+ -- compile Nat-likes as bigints
421+ -- it seems that prim applications needn't be saturated
422+ | Just LikeS <- isLikeNat ist n
423+ -> irTerm top vs env $
424+ App Complete
425+ (P Ref (sUN "prim__addBigInt") Erased)
426+ (Constant $ BI 1)
427+
408428 -- not a newtype, just apply to a constructor
409429 | otherwise
410430 -> padLams . applyToNames <$> buildApp (LV n) argsPruned
@@ -501,6 +521,58 @@ irTerm top vs env (TType _) = return LNothing
501521irTerm top vs env Erased = return LNothing
502522irTerm top vs env Impossible = return LNothing
503523
524+data LikeNat = LikeZ | LikeS
525+
526+isLikeNat :: IState -> Name -> Maybe LikeNat
527+isLikeNat ist cn
528+ -- Nat itself is special-cased in Idris/DataOpts.hs,
529+ -- which will have already happened at this point.
530+
531+ -- If the optimisation is disabled then nothing looks like Nat.
532+ | GeneralisedNatHack `notElem` opt_optimise (idris_options ist)
533+ = Nothing
534+
535+ | Just cTy <- lookupTyExact cn $ tt_ctxt ist
536+ , (P TCon{} tyN _, _) <- unApply $ getRetTy cTy
537+ , Just (z, s) <- natLikeCtors tyN
538+ = if | cn == z -> Just LikeZ
539+ | cn == s -> Just LikeS
540+ | otherwise -> error $ "isLikeNat: constructor not found in its own family: " ++ show (cn, tyN)
541+
542+ | otherwise = Nothing
543+ where
544+ natLikeCtors :: Name -> Maybe (Name, Name)
545+ natLikeCtors tyN = case lookupCtxtExact tyN $ idris_datatypes ist of
546+ Just TI{con_names = [z, s]}
547+ | 0 <- getUsedCount z
548+ , looksLikeS s
549+ -> Just (z, s)
550+
551+ Just TI{con_names = [s, z]}
552+ | 0 <- getUsedCount z
553+ , looksLikeS s
554+ -> Just (z, s)
555+
556+ _ -> Nothing
557+
558+ getUsedCount :: Name -> Int
559+ getUsedCount n =
560+ maybe 0 (length . usedpos)
561+ $ lookupCtxtExact n
562+ $ idris_callgraph ist
563+
564+ looksLikeS :: Name -> Bool
565+ looksLikeS cn
566+ | Just [(i, _)] <- fmap usedpos $ lookupCtxtExact cn $ idris_callgraph ist
567+ , Just cTy <- lookupTyExact cn $ tt_ctxt ist
568+ , [recTy] <- [recTy | (j, (_n, recTy)) <- zip [0..] (getArgTys cTy), j == i]
569+ , (P TCon{} recTyN _, _) <- unApply recTy
570+ , (P TCon{} tyN _, _) <- unApply $ getRetTy cTy
571+ , recTyN == tyN
572+ = True
573+
574+ | otherwise = False
575+
504576doForeign :: Vars -> [Name] -> [Term] -> Idris LExp
505577doForeign vs env (ret : fname : world : args)
506578 = do args' <- mapM splitArg args
@@ -616,18 +688,41 @@ irSC top vs (Case up n alts@[ConCase cn a ns sc, DefaultCase sc']) = do
616688 detag <- fgetState (opt_detaggable . ist_optimisation cn)
617689 if detag
618690 then irSC top vs (Case up n [ConCase cn a ns sc])
619− else LCase up (LV n) <$> mapM (irAlt top vs (LV n)) alts
620−
621−irSC top vs sc@(Case up n alts) = do
622− -- check that neither alternative needs the newtype optimisation,
623− -- see comment above
624− goneWrong <- or <$> mapM isDetaggable alts
625− when goneWrong
626− $ ifail ("irSC: non-trivial case-match on detaggable data: " ++ show sc)
627−
628− -- everything okay
629− LCase up (LV n) <$> mapM (irAlt top vs (LV n)) alts
691+ else do
692+ likeNat <- isLikeNat <$> getIState <*> pure cn
693+ case likeNat of
694+ -- the annoying case: LikeS is translated into a default case
695+ -- so we need to change the original DefaultCase to Z-case
696+ -- and reorder it before this one
697+ Just LikeS -> do
698+ zCase <- LConstCase (BI 0) <$> irSC top vs sc'
699+ sCase <- irAlt top vs (LV n) (ConCase cn a ns sc)
700+ return $ LCase up (LV n) [zCase, sCase]
701+
702+ -- the usual case
703+ _ -> LCase up (LV n) <$> mapM (irAlt top vs (LV n)) alts
704+
705+irSC top vs sc@(Case up n alts) = getIState >>= rhs
630706 where
707+ rhs ist
708+ | [ConCase cns as nss scs, ConCase cnz az nsz scz] <- alts
709+ , Just LikeS <- isLikeNat ist cns
710+ = do
711+ -- reorder to make the Z case come first
712+ zCase <- LConstCase (BI 0) <$> irSC top vs scz
713+ sCase <- irAlt top vs (LV n) (ConCase cns as nss scs)
714+ return $ LCase up (LV n) [zCase, sCase]
715+
716+ | otherwise = do
717+ -- check that neither alternative needs the newtype optimisation,
718+ -- see comment above
719+ goneWrong <- or <$> mapM isDetaggable alts
720+ when goneWrong
721+ $ ifail ("irSC: non-trivial case-match on detaggable data: " ++ show sc)
722+
723+ -- everything okay
724+ LCase up (LV n) <$> mapM (irAlt top vs (LV n)) alts
725+
631726 isDetaggable (ConCase cn _ _ _) = fgetState $ opt_detaggable . ist_optimisation cn
632727 isDetaggable _ = return False
633728
@@ -636,10 +731,32 @@ irSC top vs ImpossibleCase = return LNothing
636731irAlt :: Name -> Vars -> LExp -> CaseAlt -> Idris LAlt
637732
638733-- this leaves out all unused arguments of the constructor
639−irAlt top vs _ (ConCase n t args sc) = do
734+irAlt top vs tm (ConCase n t args sc) = do
735+ likeNat <- isLikeNat <$> getIState <*> pure n
640736 used <- map fst <$> fgetState (cg_usedpos . ist_callgraph n)
641737 let usedArgs = [a | (i,a) <- zip [0..] args, i `elem` used]
642− LConCase (-1) n usedArgs <$> irSC top (methodVars `M.union` vs) sc
738+
739+ case likeNat of
740+ -- like S
741+ Just LikeS -> do
742+ sc' <- irSC top vs sc
743+ case usedArgs of
744+ [pv] -> return $ LDefaultCase $ LLet pv
745+ ( LOp
746+ (LMinus (ATInt ITBig))
747+ [ tm
748+ , LConst (BI 1)
749+ ]
750+ )
751+ sc'
752+
753+ _ -> ifail $ "irAlt/LikeS: multiple used args: " ++ show usedArgs
754+
755+ -- like Z
756+ Just LikeZ -> LConstCase (BI 0) <$> irSC top vs sc
757+
758+ -- not like Nat
759+ Nothing -> LConCase (-1) n usedArgs <$> irSC top (methodVars `M.union` vs) sc
643760 where
644761 methodVars = case n of
645762 SN (ImplementationCtorN interfaceName)