@@ -5,7 +5,7 @@ Description : Coordinates the compilation process.
55License : BSD3
66Maintainer : The Idris Community.
77-}
8−{-# LANGUAGE CPP, FlexibleContexts, PatternGuards, TypeSynonymInstances, NamedFieldPuns #-}
8+{-# LANGUAGE CPP, FlexibleContexts, PatternGuards, TypeSynonymInstances, NamedFieldPuns, MultiWayIf #-}
99
1010module IRTS.Compiler(compile, generate) where
1111
@@ -255,7 +255,6 @@ type Vars = M.Map Name VarInfo
255255irTerm :: Name -> Vars -> [Name] -> Term -> Idris LExp
256256irTerm top vs env tm@(App _ f a) = do
257257 ist <- getIState
258− let (likeNatZ, likeNatS) = getNatLikeCtors ist
259258 case unApply tm of
260259 (P _ n _, args)
261260 | n `elem` map fst (idris_metavars ist) \\ primDefs
@@ -396,11 +395,11 @@ irTerm top vs env tm@(App _ f a) = do
396395 -> irTerm top vs env (head argsPruned)
397396
398397 -- compile Nat-likes as bigints
399− | n `S.member` likeNatZ
398+ | Just LikeZ <- isLikeNat ist n
400399 -> irTerm top vs env $ Constant (BI 0)
401400
402401 -- compile Nat-likes as bigints
403− | n `S.member` likeNatS
402+ | Just LikeS <- isLikeNat ist n
404403 -> irTerm top vs env $
405404 App Complete (P Ref (sUN "prim__addBigInt") Erased) (Constant $ BI 1)
406405
@@ -417,7 +416,7 @@ irTerm top vs env tm@(App _ f a) = do
417416 -> return . padLams $ \[vn] -> LApp False (LV n) [LV vn]
418417
419418 -- compile Nat-likes as bigints
420− | n `S.member` likeNatS
419+ | Just LikeS <- isLikeNat ist n
421420 -> irTerm top vs env $
422421 App Complete (P Ref (sUN "prim__addBigInt") Erased) (Constant $ BI 1)
423422
@@ -517,44 +516,47 @@ irTerm top vs env (TType _) = return LNothing
517516irTerm top vs env Erased = return LNothing
518517irTerm top vs env Impossible = return LNothing
519518
520−getNatLikeCtors :: IState -> (S.Set Name, S.Set Name)
521−getNatLikeCtors ist =
522− ( S.fromList $ map fst natLikeTypes
523− , S.fromList $ map snd natLikeTypes
524− )
519+data LikeNat = LikeZ | LikeS
520+
521+isLikeNat :: IState -> Name -> Maybe LikeNat
522+isLikeNat ist cn
523+ | Just cTy <- lookupTyExact cn $ tt_ctxt ist
524+ , (P TCon{} tyN _, _) <- unApply $ getRetTy cTy
525+ , Just (z, s) <- natLikeCtors tyN cTy
526+ = if | cn == z -> Just LikeZ
527+ | cn == s -> Just LikeS
528+ | otherwise -> Nothing
525529 where
530+ natLikeCtors :: Name -> Type -> Maybe (Name, Name)
531+ natLikeCtors tyN cTy = case lookupCtxtExact tyN $ idris_datatypes ist of
532+ Just TI{con_names = [z, s]}
533+ | 0 <- getUsedCount z
534+ , looksLikeS tyN s cTy
535+ -> Just (z, s)
536+
537+ Just TI{con_names = [s, z]}
538+ | 0 <- getUsedCount z
539+ , looksLikeS tyN s cTy
540+ -> Just (z, s)
541+
542+ _ -> Nothing
543+
526544 getUsedCount :: Name -> Int
527545 getUsedCount n =
528546 maybe 0 (length . usedpos)
529547 $ lookupCtxtExact n
530548 $ idris_callgraph ist
531549
532− isRecursive :: Name -> Name -> Bool
533− isRecursive tyN cn
534− | Just (Bind _ Pi{binderTy = P _ argTyN _} _) <- lookupTyExact cn $ tt_ctxt ist
535− , argTyN == tyN
550+ looksLikeS :: Name -> Name -> Type -> Bool
551+ looksLikeS tyN cn cTy
552+ | Just [(i, _)] <- fmap usedpos $ lookupCtxtExact cn $ idris_callgraph ist
553+ , [recTy] <- [recTy | (j, (_n, recTy)) <- zip [0..] (getArgTys cTy), j == i]
554+ , (P TCon{} recTyN _, _) <- unApply recTy
555+ , recTyN == tyN
536556 = True
537557
538558 | otherwise = False
539559
540− natLikeTypes :: [(Name, Name)]
541− natLikeTypes = do
542− (typeName, TI{con_names}) <- toAlist $ idris_datatypes ist
543− case con_names of
544− [z, s]
545− | getUsedCount z == 0
546− , getUsedCount s == 1
547− , isRecursive typeName s
548− -> [(z, s)]
549−
550− [s, z]
551− | getUsedCount z == 0
552− , getUsedCount s == 1
553− , isRecursive typeName s
554− -> [(z, s)]
555−
556− _ -> [] -- not a 2-constructor family
557−
558560doForeign :: Vars -> [Name] -> [Term] -> Idris LExp
559561doForeign vs env (ret : fname : world : args)
560562 = do args' <- mapM splitArg args
@@ -690,10 +692,32 @@ irSC top vs ImpossibleCase = return LNothing
690692irAlt :: Name -> Vars -> LExp -> CaseAlt -> Idris LAlt
691693
692694-- this leaves out all unused arguments of the constructor
693−irAlt top vs _ (ConCase n t args sc) = do
695+irAlt top vs tm (ConCase n t args sc) = do
696+ likeNat <- isLikeNat <$> getIState <*> pure n
694697 used <- map fst <$> fgetState (cg_usedpos . ist_callgraph n)
695698 let usedArgs = [a | (i,a) <- zip [0..] args, i `elem` used]
696− LConCase (-1) n usedArgs <$> irSC top (methodVars `M.union` vs) sc
699+
700+ case likeNat of
701+ -- like S
702+ Just LikeS -> do
703+ sc' <- irSC top vs sc
704+ case usedArgs of
705+ [pv] -> return $ LDefaultCase $ LLet pv
706+ ( LOp
707+ (LMinus (ATInt ITBig))
708+ [ tm
709+ , LConst (BI 1)
710+ ]
711+ )
712+ sc'
713+
714+ _ -> ifail $ "irAlt/LikeS: multiple used args: " ++ show usedArgs
715+
716+ -- like Z
717+ Just LikeZ -> LConstCase (BI 0) <$> irSC top vs sc
718+
719+ -- not like Nat
720+ Nothing -> LConCase (-1) n usedArgs <$> irSC top (methodVars `M.union` vs) sc
697721 where
698722 methodVars = case n of
699723 SN (ImplementationCtorN interfaceName)