<commit_before>{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
module IRTS.CodeGenCLaSH
  ( codeGenCLaSH
  , createBindingsCLaSH
  )
where

-- External imports
import           Control.Arrow                (first)
import           Control.Applicative          ((<$>),(<*>),pure)
import           Control.Monad.Trans          (MonadTrans)
import           Control.Monad.Trans.Error    (ErrorT(..))
import           Control.Monad.Trans.Maybe    (MaybeT(..),runMaybeT)
import           Control.Monad.Reader         (Reader)
import qualified Control.Monad.Reader         as Reader
import           Control.Monad.State          (StateT,lift,liftIO)
import qualified Control.Monad.State.Lazy     as State
import           Data.ByteString.Lazy.Char8   (pack)
import           Data.Either                  (lefts)
import           Data.HashMap.Lazy            (HashMap)
import qualified Data.HashMap.Lazy            as HashMap
import           Data.Maybe                   (catMaybes,fromMaybe)
import           Control.Lens                 ((%=),_1,_2,_3,makeLenses,view,(^.))
import qualified Unbound.LocallyNameless.Name as U
import           Unbound.LocallyNameless.Ops  (unsafeUnbind)
import           Unbound.LocallyNameless      (Bind,Rep,bind,embed,makeName,name2String,name2Integer,rebind,runFreshM,runFreshMT,string2Name,unembed)

import Debug.Trace

-- CLaSH-specific import
import qualified CLaSH.Core.DataCon as C
import qualified CLaSH.Core.Literal as C
import qualified CLaSH.Core.Pretty  as C
import qualified CLaSH.Core.Prim    as C
import qualified CLaSH.Core.Term    as C
import qualified CLaSH.Core.TyCon   as C
import qualified CLaSH.Core.Type    as C
import qualified CLaSH.Core.TysPrim as C
import qualified CLaSH.Core.Util    as C
import qualified CLaSH.Core.Var     as C
import qualified CLaSH.Rewrite.Util as C

import CLaSH.Driver
import CLaSH.Driver.Types
import CLaSH.Netlist.Util (coreTypeToHWType)
import CLaSH.Netlist.Types (HWType(..))
import CLaSH.Primitives.Types as C
import CLaSH.Primitives.Util
import CLaSH.Rewrite.Types (DebugLevel(..))
import CLaSH.Util (MonadUnique(..),traceIf)

-- Local imports
import Core.CaseTree
import Core.Evaluate
import Core.TT

import Idris.AbsSyntax
import Idris.UnusedArgs

import Debug.Trace

type MkTyDConState = (HashMap C.TyConName C.TyCon, HashMap C.TyConName [C.DataCon], HashMap C.DcName C.DataCon)
type R    = Reader MkTyDConState
type SR a = StateT MkTyDConState R a
type CoreAlt = Bind C.Pat C.Term

instance Monad m => MonadUnique (StateT Int m) where
  getUniqueM = do
    supply <- State.get
    State.modify (+1)
    return supply

codeGenCLaSH :: BindingMap
             -> PrimMap
             -> IO ()
codeGenCLaSH bindingMap primMap =
  generateVHDL bindingMap HashMap.empty HashMap.empty primMap idrisTypeToHWType DebugNone

createBindingsCLaSH :: IState
                    -> Term
                    -> [Name]
                    -> PrimMap
                    -> BindingMap
createBindingsCLaSH iState tm used primMap =
  let ctxtList  = ctxtAlist (tt_ctxt iState)
      dcTyCons  = filter (isCon . snd) ctxtList
      decs      = filter ((`elem` used) . fst) ctxtList
      dcTcState = makeAllTyDataCons (map defToTyDataCon dcTyCons)
      terms     = map (defToTerm primMap dcTcState) decs
  in traceIf False (unlines $ map (\(n,d) -> C.showDoc n ++ ": " ++ C.showDoc (C.dcType d)) $ HashMap.toList (dcTcState ^. _3))
     (HashMap.fromList (catMaybes terms))

isCon :: Def -> Bool
isCon (TyDecl _ _) = True
isCon _ = False

defToTyDataCon :: (Name,Def)
               -> (Name,NameType,Type)
defToTyDataCon (n,TyDecl nt ty) = (n,nt,ty)
defToTyDataCon nd = error $ "Not a TyDec" ++ show nd

isTyConDef :: (Name,NameType,Type)
           -> Bool
isTyConDef (_,TCon _ _,_) = True
isTyConDef _              = False

isDConDef :: (Name,NameType,Type)
          -> Bool
isDConDef (_,DCon _ _,_) = True
isDConDef _              = False

makeAllTyDataCons :: [(Name,NameType,Type)]
                  -> MkTyDConState
makeAllTyDataCons tyDecls =
  let s = Reader.runReader (State.execStateT
                              (do mapM_ makeTyCon (filter isTyConDef tyDecls)
                                  mapM_ makeDataCon (filter isDConDef tyDecls)
                              )
                              (HashMap.empty,HashMap.empty,HashMap.empty)
                           ) s
  in  s

makeDataCon,makeTyCon :: (Name,NameType,Type) -> SR ()

makeDataCon k@(n,DCon t a,ty) = do
  dcTyM <- lift $ runMaybeT $ toCoreType False [] [] ty
  dcTy  <- case dcTyM of
             (Just t) -> return t
             Nothing  -> error $ "makeDataCon: " ++ showTerm ty
  let (args,resTy)     = splitFunForallTy dcTy
      (C.TyConApp _ l) = C.tyView resTy
      (tvs,ids)        = splitAt (length l) args
      tyVars           = lefts tvs
      argTys           = map (either (unembed . C.varKind) id) ids
      dcName = toCoreName n
      dc = C.MkData { C.dcName       = dcName
                    , C.dcTag        = t
                    , C.dcType       = dcTy
                    , C.dcArgTys     = argTys
                    , C.dcUnivTyVars = map C.varName tyVars
                    , C.dcExtTyVars  = []
                    }

      tcName   = toCoreName n
      liftedDc = C.AlgTyCon { C.tyConName = tcName
                            , C.tyConKind = dcTy
                            , C.tyConArity = (length tyVars + length argTys)
                            , C.algTcRhs = C.DataTyCon []
                            , C.isDictTyCon = False
                            }

  case (C.splitTyConAppM resTy) of
    Just (tc,_) -> do
      _1 %= HashMap.insert tcName liftedDc
      _2 %= HashMap.insertWith (++) (C.tyConName tc) [dc]
      _3 %= HashMap.insert dcName dc
    Nothing -> error $ "Huh?: " ++ show ty

makeTyCon (n,TCon t a,ty) = do
  let tcName = toCoreName n
  dcs  <- lift $ tcDataCons tcName
  tcKiM <- lift $ runMaybeT $ toCoreType False [] [] ty
  tcKi <- case tcKiM of
             (Just k) -> return k
             Nothing  -> error $ "makeTyCon: " ++ showTerm ty
  let tc = C.AlgTyCon { C.tyConName   = tcName
                      , C.tyConKind   = tcKi
                      , C.tyConArity  = a
                      , C.algTcRhs    = C.DataTyCon dcs
                      , C.isDictTyCon = False
                      }
  _1 %= HashMap.insert tcName tc

tcDataCons :: C.TyConName
           -> R [C.DataCon]
tcDataCons tc = fmap ( fromMaybe (error $ "No DataCons found for: " ++ show tc)
                     . HashMap.lookup tc
                     )
              $ view _2

toTyCon :: C.TyConName
        -> R C.TyCon
toTyCon tcN = fmap ( fromMaybe (error $ "No TyCon named: " ++ name2String tcN)
                   . HashMap.lookup tcN
                   )
            $ view _1

toDataCon :: C.DcName
          -> R C.DataCon
toDataCon dcN = fmap ( fromMaybe (error $ "No DataCon named: " ++ name2String dcN)
                   . HashMap.lookup dcN
                   )
              $ view _3

toCoreName :: Rep a => Name -> U.Name a
toCoreName (UN s)   = string2Name s
toCoreName (NS n s) = prefixName (showSep "." (reverse s) ++ ".") (toCoreName n)
toCoreName (MN i s) = makeName s (toInteger i)
toCoreName NErased  = string2Name "___ERASED___"

prefixName :: String -> U.Name a -> U.Name a
prefixName p (U.Nm r (s,i)) = U.Nm r (p++s,i)
prefixName _ n              = n

replaceNameIndex :: Integer -> U.Name a -> U.Name a
replaceNameIndex i (U.Nm r (s,_)) = U.Nm r (s,i)
replaceNameIndex _ n              = n

toCoreTypeT :: Bool -> [(C.TyName,(C.Kind,Type))] -> [(Name,Either C.Id C.TyVar)] -> Type -> MaybeT R C.Type
toCoreTypeT refKI ns bndrs t = do
  tC <- toCoreType refKI ns bndrs t
  return $! toCoreTypeT' tC

toCoreTypeT' :: C.Type -> C.Type
toCoreTypeT' tC = case (C.tyView tC) of
  (C.TyConApp tc args) -> case (name2String $ C.tyConName tc) of
      "Builtins.fromInteger" -> (last args)
      "@Builtins.Num$[Nat].0.Builtins.#!fromInteger" -> (last args)
      "@Builtins.Num$[Nat].0.Builtins.#!fromInteger.0.#fromInteger'" -> (last args)
      _ -> C.mkTyConApp tc (map toCoreTypeT' args)
  (C.FunTy arg res) -> C.mkFunTy (toCoreTypeT' arg) (toCoreTypeT' res)
  _ -> tC

toCoreType :: Bool -> [(C.TyName,(C.Kind,Type))] -> [(Name,Either C.Id C.TyVar)] -> Type -> MaybeT R C.Type
toCoreType refKI ns bndrs t@(Bind n (Pi ki) tt)
  | looksLikeKind refKI ns ki
  = do
    kiC <- toCoreKind ns ki
    let tvN = toCoreName n
        tv  = C.TyVar tvN (embed kiC)
    C.ForAllTy <$> (bind tv <$> toCoreType refKI ((tvN,(kiC,ki)):ns) bndrs tt)

toCoreType refKI ns bndrs t@(Bind n (Pi arg) res) = do
  let tvN = toCoreName n
  C.mkFunTy <$> toCoreType refKI ns bndrs arg <*> toCoreType refKI ((tvN,(error $ "unexpected kind: " ++ show arg,arg)):ns) bndrs res

toCoreType refKI ns bndrs t@(P (TCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreType refKI ns bndrs t@(P (DCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreType refKI ns bndrs t@(P Ref n _)        = return (C.mkTyConTy $ C.mkPrimTyCon (toCoreName n) C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(P Bound n _)      = case lookup n bndrs of
                                                   Just (Left id_) -> return $! C.mkTyVarTy (unembed $ C.varType id_) (U.translate $ C.varName id_)
                                                   Just (Right tv) -> return $! C.mkTyVarTy (unembed $ C.varKind tv) (C.varName tv)
                                                   Nothing -> error ("toCoreType Bound notfound: " ++ showTerm t ++ show bndrs)
toCoreType refKI ns bndrs t@(App ty1 ty2) = C.AppTy <$> toCoreType refKI ns bndrs ty1 <*> toCoreType refKI ns bndrs ty2
toCoreType refKI ns bndrs t@(V i) = let (tyN,(tyK,_)) = maybe (error ("Index(toCoreType): " ++ show i ++ " not found in: " ++ show ns)) id (safeIndex ns i)
                                    in return $! C.mkTyVarTy tyK tyN
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITBig))) = return C.intPrimTy -- return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "GHC.Integer.Type.Integer") C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITChar))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "Char") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITNative))) = return C.intPrimTy -- return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "GHC.Integer.Type.Integer") C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt (ITFixed _)))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name $ show t) C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt (ITVec _ _)))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name $ show t) C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType ATFloat)) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "Float") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant (Str s))               = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name s) C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant StrType) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "String") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant PtrType) = return C.voidPrimTy
toCoreType refKI ns bndrs t@(Constant (BI i)) = return $! C.LitTy (C.NumTy $! fromInteger i)
toCoreType refKI ns bndrs t@(Constant c) = error $ "toCoreType: Constant(" ++ showConstant c ++ ")"
toCoreType refKI ns bndrs t@(TType _) = return C.liftedTypeKind
toCoreType refKI _ _ t = error $ "toCoreType: " ++ show t


looksLikeKind :: Bool
              -> [(C.TyName,(C.Kind,Type))]
              -> Type
              -> Bool
looksLikeKind _ ns (TType _)           = True
looksLikeKind refKI ns (Bind n (Pi ki) tt) = looksLikeKind refKI ns ki && looksLikeKind refKI ((toCoreName n,(error "lookslikekind",ki)):ns) tt
looksLikeKind _ ns (P (TCon _ _) n@(NS (UN "Nat") ["Nat","Prelude"]) _) = True
-- looksLikeKind True ns (V i) = maybe (error ("Index: " ++ show i ++ " not found in: " ++ show ns)) (looksLikeKind True ns . snd . snd) (safeIndex ns i)
looksLikeKind _ _ k                   = False

toCoreKind :: [(C.TyName,(C.Kind,Type))] -> Type -> MaybeT R C.Kind
toCoreKind ns t@(TType _) = return C.liftedTypeKind
toCoreKind ns t@(Bind n (Pi arg) res) = do
  let tvN = toCoreName n
  C.mkFunTy <$> toCoreKind ns arg <*> toCoreKind ((tvN,error $ "unexpected box: " ++ show arg):ns) res
toCoreKind ns t@(P (TCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreKind ns t@(V i) = toCoreKind ns (snd $ snd (ns !! i))
toCoreKind ns t = error ("toCoreKind: " ++ show t ++ "\n" ++ showTerm t)

defToTerm :: PrimMap
          -> MkTyDConState
          -> (Name,Def)
          -> Maybe (C.TmName,(String,(C.Type,C.Term)))
defToTerm primMap _ (n,_)
  | HashMap.member (pack $ name2String $ (toCoreName n :: C.TmName)) primMap = Nothing
defToTerm primMap s (n,CaseOp _ ty _ _ ns sc _ _) = traceIf False ("== " ++ show n ++ " ==:\nTy: " ++ show ty ++ "\nBndrs: " ++ show ns ++ "\nTm:\n" ++ show sc) $ flip Reader.runReader s $ do
  let tmName   = toCoreName n
      tmString = show n
  tyC <- fmap (maybe (error $ "defToTerm: " ++ showTerm ty) id) $ runMaybeT $ toCoreTypeT False [] [] ty
  let (args,resTy) = splitFunForallTy tyC
      bndrs        = zipWith (\n ds -> case ds of { Left (C.TyVar n' t) -> Right (C.TyVar n' t)
                                                  ; Right t             -> Left  (C.Id (toCoreName n) (embed t))
                                                  }) ns args
  result <- runMaybeT $ do { scC <- scToTerm primMap [] (zip ns bndrs) sc
                           ; let termC  = C.mkAbstraction scC bndrs
                           ; let res    = (tmName,(tmString,(tyC,termC)))
                           ; let resTy  = runFreshM $ C.termType termC
                           ; traceIf False ("CoreTy: " ++ C.showDoc tyC ++ "\nCoreTerm:\n" ++ C.showDoc termC) $
                               traceIf (tyC /= resTy) ("TYPES DIFFER for " ++ tmString ++ ":\nExpected: " ++ C.showDoc tyC ++ "\nCalculated: " ++ C.showDoc resTy) $
                                 return res
                           }
  return result
  -- return Nothing

defToTerm _ _ _ = Nothing

scToTerm :: PrimMap
         -> [(C.TyName,(C.Kind,Type))]
         -> [(Name,Either C.Id C.TyVar)]
         -> SC<commit_msg>Implement blockSinking (part of lambda-dropping)<commit_after>{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
module IRTS.CodeGenCLaSH
  ( codeGenCLaSH
  , createBindingsCLaSH
  )
where

-- External imports
import           Control.Arrow                (first,second)
import           Control.Applicative          ((<$>),(<*>),pure)
import           Control.Monad.Trans          (MonadTrans)
import           Control.Monad.Trans.Error    (ErrorT(..))
import           Control.Monad.Trans.Maybe    (MaybeT(..),runMaybeT)
import           Control.Monad.Reader         (Reader)
import qualified Control.Monad.Reader         as Reader
import           Control.Monad.State          (StateT,lift,liftIO)
import qualified Control.Monad.State.Lazy     as State
import           Data.ByteString.Lazy.Char8   (pack)
import           Data.Either                  (lefts)
import           Data.Graph.Inductive         (Gr,LNode,lsuc,mkGraph,iDom)
import           Data.HashMap.Lazy            (HashMap)
import qualified Data.HashMap.Lazy            as HashMap
import qualified Data.List                    as List
import           Data.Maybe                   (catMaybes,fromMaybe,listToMaybe)
import           Control.Lens                 ((%=),_1,_2,_3,makeLenses,view,(^.))
import qualified Unbound.LocallyNameless.Name as U
import           Unbound.LocallyNameless.Ops  (unsafeUnbind)
import           Unbound.LocallyNameless      (Bind,Rep,bind,embed,makeName,name2String,name2Integer,rebind,rec,runFreshM,runFreshMT,string2Name,unembed)

import Debug.Trace

-- CLaSH-specific import
import qualified CLaSH.Core.DataCon as C
import qualified CLaSH.Core.Literal as C
import qualified CLaSH.Core.Pretty  as C
import qualified CLaSH.Core.Prim    as C
import qualified CLaSH.Core.Term    as C
import qualified CLaSH.Core.TyCon   as C
import qualified CLaSH.Core.Type    as C
import qualified CLaSH.Core.TysPrim as C
import qualified CLaSH.Core.Util    as C
import qualified CLaSH.Core.Var     as C
import qualified CLaSH.Rewrite.Util as C

import CLaSH.Driver
import CLaSH.Driver.Types
import CLaSH.Netlist.Util (coreTypeToHWType)
import CLaSH.Netlist.Types (HWType(..))
import CLaSH.Normalize.Util (callGraph,recursiveComponents)
import CLaSH.Primitives.Types as C
import CLaSH.Primitives.Util
import CLaSH.Rewrite.Types (DebugLevel(..))
import CLaSH.Util (MonadUnique(..),traceIf)

-- Local imports
import Core.CaseTree
import Core.Evaluate
import Core.TT

import Idris.AbsSyntax
import Idris.UnusedArgs

import Debug.Trace

type MkTyDConState = (HashMap C.TyConName C.TyCon, HashMap C.TyConName [C.DataCon], HashMap C.DcName C.DataCon)
type R    = Reader MkTyDConState
type SR a = StateT MkTyDConState R a
type CoreAlt = Bind C.Pat C.Term

instance Monad m => MonadUnique (StateT Int m) where
  getUniqueM = do
    supply <- State.get
    State.modify (+1)
    return supply

codeGenCLaSH :: BindingMap
             -> PrimMap
             -> IO ()
codeGenCLaSH bindingMap primMap =
  generateVHDL bindingMap HashMap.empty HashMap.empty primMap idrisTypeToHWType DebugFinal

createBindingsCLaSH :: IState
                    -> Term
                    -> [Name]
                    -> PrimMap
                    -> BindingMap
createBindingsCLaSH iState tm used primMap =
  let ctxtList  = ctxtAlist (tt_ctxt iState)
      dcTyCons  = filter (isCon . snd) ctxtList
      decs      = filter ((`elem` used) . fst) ctxtList
      dcTcState = makeAllTyDataCons (map defToTyDataCon dcTyCons)
      terms     = catMaybes $ map (defToTerm primMap dcTcState) decs
      terms'    = termPrep terms
  in traceIf False (unlines $ map (\(n,d) -> C.showDoc n ++ ": " ++ C.showDoc (C.dcType d)) $ HashMap.toList (dcTcState ^. _3))
     terms'


termPrep :: [(C.TmName,(C.Type,C.Term))]
             -> BindingMap
termPrep bndrs = maybe (HashMap.empty) id $ do
  (topEntity,_) <- listToMaybe $ filter (List.isSuffixOf "topEntity" . name2String . fst) bndrs
  let depGraph = callGraph [] (HashMap.fromList $ map (second snd) bndrs) topEntity
      used     = HashMap.fromList depGraph
      bndrs'   = HashMap.fromList $ filter ((`HashMap.member` used) . fst) bndrs
      rcs      = recursiveComponents depGraph
      dropped  = map (lambdaDrop bndrs' used) rcs
      bndrs''  = foldr (\(k,v) b -> HashMap.insert k v b) bndrs' dropped
  return bndrs''

lambdaDrop :: BindingMap
           -> HashMap C.TmName [C.TmName]
           -> [C.TmName]
           -> (C.TmName,(C.Type,C.Term))
lambdaDrop bndrs cfg cyc@(root:_) = block
  where
    doms  = dominator cfg cyc
    block = blockSink bndrs doms (0,root)

blockSink :: BindingMap
          -> Gr C.TmName C.TmName
          -> LNode C.TmName
          -> (C.TmName,(C.Type,C.Term))
blockSink bndrs doms (nId,tmName) = (tmName,(ty,newTm))
  where
    (ty,tm) = bndrs HashMap.! tmName
    sucTm   = lsuc doms nId
    tmS     = map (blockSink bndrs doms) sucTm
    bnds    = map (\(tN,(ty',tm')) -> (C.Id tN (embed ty'),embed tm')) tmS
    (tmArgs,tmExpr) = runFreshM $ C.collectBndrs tm
    newTm   = case sucTm of
                [] -> tm
                _  -> C.mkAbstraction (C.Letrec $ bind (rec bnds) tmExpr) tmArgs


dominator :: HashMap C.TmName [C.TmName]
          -> [C.TmName]
          -> Gr C.TmName C.TmName
dominator cfg cyc@(root:others) = mkGraph nodes (map (\(e,b) -> (b,e,nodesM HashMap.! e)) doms)
  where
    nodes    = zip [0..] cyc
    nodesM   = HashMap.fromList nodes
    nodesI   = HashMap.fromList $ zip cyc [0..]
    cycEdges = HashMap.map ( map (nodesI HashMap.!)
                           . filter (`elem` cyc)
                           )
             $ HashMap.filterWithKey (\k v -> k `elem` cyc) cfg
    edges    = concatMap (\(i,n) -> zip3 (repeat i) (cycEdges HashMap.! n) (repeat ())
                         ) nodes
    graph    = mkGraph nodes edges :: Gr C.TmName ()
    doms     = iDom graph 0

isCon :: Def -> Bool
isCon (TyDecl _ _) = True
isCon _ = False

defToTyDataCon :: (Name,Def)
               -> (Name,NameType,Type)
defToTyDataCon (n,TyDecl nt ty) = (n,nt,ty)
defToTyDataCon nd = error $ "Not a TyDec" ++ show nd

isTyConDef :: (Name,NameType,Type)
           -> Bool
isTyConDef (_,TCon _ _,_) = True
isTyConDef _              = False

isDConDef :: (Name,NameType,Type)
          -> Bool
isDConDef (_,DCon _ _,_) = True
isDConDef _              = False

makeAllTyDataCons :: [(Name,NameType,Type)]
                  -> MkTyDConState
makeAllTyDataCons tyDecls =
  let s = Reader.runReader (State.execStateT
                              (do mapM_ makeTyCon (filter isTyConDef tyDecls)
                                  mapM_ makeDataCon (filter isDConDef tyDecls)
                              )
                              (HashMap.empty,HashMap.empty,HashMap.empty)
                           ) s
  in  s

makeDataCon,makeTyCon :: (Name,NameType,Type) -> SR ()

makeDataCon k@(n,DCon t a,ty) = do
  dcTyM <- lift $ runMaybeT $ toCoreType False [] [] ty
  dcTy  <- case dcTyM of
             (Just t) -> return t
             Nothing  -> error $ "makeDataCon: " ++ showTerm ty
  let (args,resTy)     = splitFunForallTy dcTy
      (C.TyConApp _ l) = C.tyView resTy
      (tvs,ids)        = splitAt (length l) args
      tyVars           = lefts tvs
      argTys           = map (either (unembed . C.varKind) id) ids
      dcName = toCoreName n
      dc = C.MkData { C.dcName       = dcName
                    , C.dcTag        = t
                    , C.dcType       = dcTy
                    , C.dcArgTys     = argTys
                    , C.dcUnivTyVars = map C.varName tyVars
                    , C.dcExtTyVars  = []
                    }

      tcName   = toCoreName n
      liftedDc = C.AlgTyCon { C.tyConName = tcName
                            , C.tyConKind = dcTy
                            , C.tyConArity = (length tyVars + length argTys)
                            , C.algTcRhs = C.DataTyCon []
                            , C.isDictTyCon = False
                            }

  case (C.splitTyConAppM resTy) of
    Just (tc,_) -> do
      _1 %= HashMap.insert tcName liftedDc
      _2 %= HashMap.insertWith (++) (C.tyConName tc) [dc]
      _3 %= HashMap.insert dcName dc
    Nothing -> error $ "Huh?: " ++ show ty

makeTyCon (n,TCon t a,ty) = do
  let tcName = toCoreName n
  dcs  <- lift $ tcDataCons tcName
  tcKiM <- lift $ runMaybeT $ toCoreType False [] [] ty
  tcKi <- case tcKiM of
             (Just k) -> return k
             Nothing  -> error $ "makeTyCon: " ++ showTerm ty
  let tc = C.AlgTyCon { C.tyConName   = tcName
                      , C.tyConKind   = tcKi
                      , C.tyConArity  = a
                      , C.algTcRhs    = C.DataTyCon dcs
                      , C.isDictTyCon = False
                      }
  _1 %= HashMap.insert tcName tc

tcDataCons :: C.TyConName
           -> R [C.DataCon]
tcDataCons tc = fmap ( fromMaybe (error $ "No DataCons found for: " ++ show tc)
                     . HashMap.lookup tc
                     )
              $ view _2

toTyCon :: C.TyConName
        -> R C.TyCon
toTyCon tcN = fmap ( fromMaybe (error $ "No TyCon named: " ++ name2String tcN)
                   . HashMap.lookup tcN
                   )
            $ view _1

toDataCon :: C.DcName
          -> R C.DataCon
toDataCon dcN = fmap ( fromMaybe (error $ "No DataCon named: " ++ name2String dcN)
                   . HashMap.lookup dcN
                   )
              $ view _3

toCoreName :: Rep a => Name -> U.Name a
toCoreName (UN s)   = string2Name s
toCoreName (NS n s) = prefixName (showSep "." (reverse s) ++ ".") (toCoreName n)
toCoreName (MN i s) = makeName s (toInteger i)
toCoreName NErased  = string2Name "___ERASED___"

prefixName :: String -> U.Name a -> U.Name a
prefixName p (U.Nm r (s,i)) = U.Nm r (p++s,i)
prefixName _ n              = n

replaceNameIndex :: Integer -> U.Name a -> U.Name a
replaceNameIndex i (U.Nm r (s,_)) = U.Nm r (s,i)
replaceNameIndex _ n              = n

toCoreTypeT :: Bool -> [(C.TyName,(C.Kind,Type))] -> [(Name,Either C.Id C.TyVar)] -> Type -> MaybeT R C.Type
toCoreTypeT refKI ns bndrs t = do
  tC <- toCoreType refKI ns bndrs t
  return $! toCoreTypeT' tC

toCoreTypeT' :: C.Type -> C.Type
toCoreTypeT' tC = case (C.tyView tC) of
  (C.TyConApp tc args) -> case (name2String $ C.tyConName tc) of
      "Builtins.fromInteger" -> (last args)
      "@Builtins.Num$[Nat].0.Builtins.#!fromInteger" -> (last args)
      "@Builtins.Num$[Nat].0.Builtins.#!fromInteger.0.#fromInteger'" -> (last args)
      _ -> C.mkTyConApp tc (map toCoreTypeT' args)
  (C.FunTy arg res) -> C.mkFunTy (toCoreTypeT' arg) (toCoreTypeT' res)
  _ -> tC

toCoreType :: Bool -> [(C.TyName,(C.Kind,Type))] -> [(Name,Either C.Id C.TyVar)] -> Type -> MaybeT R C.Type
toCoreType refKI ns bndrs t@(Bind n (Pi ki) tt)
  | looksLikeKind refKI ns ki
  = do
    kiC <- toCoreKind ns ki
    let tvN = toCoreName n
        tv  = C.TyVar tvN (embed kiC)
    C.ForAllTy <$> (bind tv <$> toCoreType refKI ((tvN,(kiC,ki)):ns) bndrs tt)

toCoreType refKI ns bndrs t@(Bind n (Pi arg) res) = do
  let tvN = toCoreName n
  C.mkFunTy <$> toCoreType refKI ns bndrs arg <*> toCoreType refKI ((tvN,(error $ "unexpected kind: " ++ show arg,arg)):ns) bndrs res

toCoreType refKI ns bndrs t@(P (TCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreType refKI ns bndrs t@(P (DCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreType refKI ns bndrs t@(P Ref n _)        = return (C.mkTyConTy $ C.mkPrimTyCon (toCoreName n) C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(P Bound n _)      = case lookup n bndrs of
                                                   Just (Left id_) -> return $! C.mkTyVarTy (unembed $ C.varType id_) (U.translate $ C.varName id_)
                                                   Just (Right tv) -> return $! C.mkTyVarTy (unembed $ C.varKind tv) (C.varName tv)
                                                   Nothing -> error ("toCoreType Bound notfound: " ++ showTerm t ++ show bndrs)
toCoreType refKI ns bndrs t@(App ty1 ty2) = C.AppTy <$> toCoreType refKI ns bndrs ty1 <*> toCoreType refKI ns bndrs ty2
toCoreType refKI ns bndrs t@(V i) = let (tyN,(tyK,_)) = maybe (error ("Index(toCoreType): " ++ show i ++ " not found in: " ++ show ns)) id (safeIndex ns i)
                                    in return $! C.mkTyVarTy tyK tyN
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITBig))) = return C.intPrimTy -- return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "GHC.Integer.Type.Integer") C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITChar))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "Char") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt ITNative))) = return C.intPrimTy -- return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "GHC.Integer.Type.Integer") C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt (ITFixed _)))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name $ show t) C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType (ATInt (ITVec _ _)))) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name $ show t) C.liftedTypeKind 0 C.IntRep)
toCoreType refKI ns bndrs t@(Constant (AType ATFloat)) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "Float") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant (Str s))               = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name s) C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant StrType) = return (C.mkTyConTy $ C.mkPrimTyCon (string2Name "String") C.liftedTypeKind 0 C.VoidRep)
toCoreType refKI ns bndrs t@(Constant PtrType) = return C.voidPrimTy
toCoreType refKI ns bndrs t@(Constant (BI i)) = return $! C.LitTy (C.NumTy $! fromInteger i)
toCoreType refKI ns bndrs t@(Constant c) = error $ "toCoreType: Constant(" ++ showConstant c ++ ")"
toCoreType refKI ns bndrs t@(TType _) = return C.liftedTypeKind
toCoreType refKI _ _ t = error $ "toCoreType: " ++ show t


looksLikeKind :: Bool
              -> [(C.TyName,(C.Kind,Type))]
              -> Type
              -> Bool
looksLikeKind _ ns (TType _)           = True
looksLikeKind refKI ns (Bind n (Pi ki) tt) = looksLikeKind refKI ns ki && looksLikeKind refKI ((toCoreName n,(error "lookslikekind",ki)):ns) tt
looksLikeKind _ ns (P (TCon _ _) n@(NS (UN "Nat") ["Nat","Prelude"]) _) = True
-- looksLikeKind True ns (V i) = maybe (error ("Index: " ++ show i ++ " not found in: " ++ show ns)) (looksLikeKind True ns . snd . snd) (safeIndex ns i)
looksLikeKind _ _ k                   = False

toCoreKind :: [(C.TyName,(C.Kind,Type))] -> Type -> MaybeT R C.Kind
toCoreKind ns t@(TType _) = return C.liftedTypeKind
toCoreKind ns t@(Bind n (Pi arg) res) = do
  let tvN = toCoreName n
  C.mkFunTy <$> toCoreKind ns arg <*> toCoreKind ((tvN,error $ "unexpected box: " ++ show arg):ns) res
toCoreKind ns t@(P (TCon _ _) n _) = C.mkTyConTy <$> lift (toTyCon (toCoreName n))
toCoreKind ns t@(V i) = toCoreKind ns (snd $ snd (ns !! i))
toCoreKind ns t = error ("toCoreKind: " ++ show t ++ "\n" ++ showTerm t)

defToTerm :: PrimMap
          -> MkTyDConState
          -> (Name,Def)
          -> Maybe (C.TmName,(C.Type,C.Term))
defToTerm primMap _ (n,_)
  | HashMap.member (pack $ name2String $ (toCoreName n :: C.TmName)) primMap = Nothing
defToTerm primMap s (n,CaseOp _ ty _ _ ns sc _ _) = traceIf False ("== " ++ show n ++ " ==:\nTy: " ++ show ty ++ "\nBndrs: " ++ show ns ++ "\nTm:\n" ++ show sc) $ flip Reader.runReader s $ do
  let tmName   = toCoreName n
      tmString = show n
  tyC <- fmap (maybe (error $ "defToTerm: " ++ showTerm ty) id) $ runMaybeT $ toCoreTypeT False [] [] ty
  let (args,resTy) = splitFunForallTy tyC
      bndrs        = zipWith (\n ds -> case ds of { Left (C.TyVar n' t) -> Right (C.TyVar n' t)
                                                  ; Right t             -> Left  (C.Id (toCoreName n) (embed t))
                                                  }) ns args
  result <- runMaybeT $ do { scC <- scToTerm primMap [] (zip ns bndrs) sc
                           ; let termC  = C.mkAbstraction scC bndrs
                           ; let res    = (tmName,(tyC,termC))
                           ; let resTy  = runFreshM $ C.termType termC
                           ; traceIf False ("CoreTy: " ++ C.showDoc tyC ++ "\nCoreTerm:\n" ++ C.showDoc termC) $
                               traceIf (tyC /= resTy) ("TYPES DIFFER for " ++ tmString ++ ":\nExpected: " ++ C.showDoc tyC ++ "\nCalculated: " ++ C.showDoc resTy) $
                                 return res
                           }
  return result
  -- return Nothing

defToTerm _ _ _ = Nothing

scToTerm :: PrimMap
         -> [(C.TyName,(C.Kind,Type))]
         -> [(Name,Either C.Id C.TyVar)]
         -> SC
{'fraction_non_alphanumeric': 0.11813846301571553, 'fraction_numerical': 0.0031367026213871907, 'mean_word_length': 2.179403683728503, 'pattern_counts': {'":': 4, '<': 52, '<?xml version=': 0, '>': 169, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 34, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4705715', 'n_tokens_mistral': 10967, 'n_tokens_neox': 10726, 'n_words': 2900}