Repository: HigherOrderCO/kind
Branch: master
Commit: 5cfff210b3ae
Files: 19
Total size: 168.7 KB

Directory structure:
gitextract_7vmntj28/

├── .gitignore
├── CHANGELOG.md
├── LICENSE
├── README.md
├── app/
│   └── Main.hs
├── cabal.project
├── kind-lang.cabal
├── main.kindc
└── src/
    ├── Kind/
    │   ├── CLI.hs
    │   ├── Check.hs
    │   ├── CompileJS.hs
    │   ├── Env.hs
    │   ├── Equal.hs
    │   ├── Parse.hs
    │   ├── Reduce.hs
    │   ├── Show.hs
    │   ├── Type.hs
    │   └── Util.hs
    └── Kind.hs

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
dist-*
*.o
*.hi
*.chi
*.chs.h
*.dyn_o
*.dyn_hi
.cabal-sandbox/
cabal.sandbox.config
*.prof
*.aux
*.hp
*.eventlog
.stack-work/
cabal.project.local*
.ghc.environment.*
.vscode/
.idea/
.DS_Store
.holefill
.tmp
.backup/
*.koder


================================================
FILE: CHANGELOG.md
================================================
# Revision history for kind2hs

## 0.1.0.0 -- YYYY-mm-dd

* First version. Released on an unsuspecting world.


================================================
FILE: LICENSE
================================================
Copyright (c) 2024 Victor Taelin

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.


================================================
FILE: README.md
================================================
# Kind

Kind is a minimal Proof Checker.

This repository is a full rewrite of Kind from the old JS implementation to
Haskell. It is being actively developed. See examples on
[KindBook](https://github.com/HigherOrderCO/KindBook).

# Usage

1. Clone and install this project

2. Use the `kind` command to check/run terms

---

We will write a proper README later (:


================================================
FILE: app/Main.hs
================================================
module Main where

import Kind (main)

main :: IO ()
main = Kind.main


================================================
FILE: cabal.project
================================================
packages: .

-- Enable -O2 optimization for all packages
package *
  optimization: 2


================================================
FILE: kind-lang.cabal
================================================
cabal-version:      3.0
name:               kind-lang
version:            0.1.0.0
license:            MIT
license-file:       LICENSE
author:             Victor Taelin
maintainer:         victor.taelin@gmail.com
category:           Language
build-type:         Simple
extra-doc-files:    CHANGELOG.md

common warnings
    ghc-options: -w

library
    import:           warnings
    exposed-modules:  Kind
                    , Kind.CLI
                    , Kind.Check
                    , Kind.CompileJS
                    , Kind.Env
                    , Kind.Equal
                    , Kind.Parse
                    , Kind.Reduce
                    , Kind.Show
                    , Kind.Type
                    , Kind.Util
    other-modules:    
    build-depends:    base ^>=4.20.0.0
                    , containers ==0.7
                    , parsec ==3.1.17.0
                    , ansi-terminal==1.1.1
                    , directory==1.3.8.3
                    , hs-highlight == 1.0.3
                    , filepath==1.5.2.0
                    , mtl==2.3.1
    hs-source-dirs:   src
    default-language: GHC2024

executable kind
    import:           warnings
    main-is:          Main.hs
    build-depends:    base ^>=4.20.0.0
                    , kind-lang
                    , ansi-terminal==1.1.1
                    , directory==1.3.8.3
                    , hs-highlight == 1.0.3
                    , filepath==1.5.2.0
                    , mtl==2.3.1
    hs-source-dirs:   app
    default-language: GHC2024


================================================
FILE: main.kindc
================================================


Bool : * = #[]{
  #true{} : Bool
  #false{} : Bool
};

Nat : * = #[]{
  #zero{} : Nat
  #succ{ pred: Nat } : Nat
};


IsTrue : ∀(b: Bool) * = λb #[b]{
  #indeed{} : (IsTrue #true{})
};

Equal : ∀(T: *) ∀(a: T) ∀(b: T) * = λT λa λb #[a b]{
  #refl{} : (Equal T a a)
};

rewrite
: ∀(T: *)
  ∀(a: T)
  ∀(b: T)
  ∀(e: (Equal T a b))
  ∀(P: ∀(x: A) *)
  ∀(x: (P a))
  (P b)
= λT λa λb λ{
  #refl: λP λx x
};

MAIN = rewrite;

//MAIN
//: (((Equal Bool) #true{}) #true{})
//= #refl{};

//Equal.rewrite<A: Type, a: A, b: A>(e: Equal<A,a,b>)<P: A -> Type>(x: P(a)): P(b)
  //case e {
    //refl: x
  //} : P(e.b)



================================================
FILE: src/Kind/CLI.hs
================================================
-- Type.hs:
-- //./Type.hs//

module Kind.CLI where

import Control.Exception (try)
import Control.Monad (forM, forM_, foldM)
import Data.List (stripPrefix, isSuffixOf, nub)
import Highlight (highlightError)
import Kind.Check
import Kind.CompileJS
import Kind.Env
import Kind.Parse
import Kind.Reduce
import Kind.Show
import Kind.Type
import Kind.Util
import System.Console.ANSI
import System.Directory (canonicalizePath, getCurrentDirectory, doesDirectoryExist, doesFileExist, getDirectoryContents)
import System.Environment (getArgs)
import System.Exit (exitWith, ExitCode(ExitSuccess, ExitFailure))
import System.FilePath (takeDirectory, (</>), takeFileName, dropExtension, isExtensionOf)
import System.IO (readFile)
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Set as S

import Debug.Trace

type FileCtx = (Book, M.Map FilePath [String], M.Map FilePath [String])
type Command = String -> FileCtx -> String -> String -> IO (Either String ())

-- main :: IO ()
-- main = ctest

main :: IO ()
main = do
  args     <- getArgs
  currPath <- getCurrentDirectory
  bookPath <- findBookDir currPath
  case bookPath of
    Nothing -> do
      putStrLn "Error: No 'book' directory found in the path."
      exitWith (ExitFailure 1)
    Just bookPath -> do
      result <- case args of
        -- ["check"]      -> runWithAll bookPath cliCheckAll
        ["run", arg]   -> runWithOne bookPath arg cliNormal
        ["check"]      -> runWithAll bookPath cliCheck
        ["check", arg] -> runWithOne bookPath arg cliCheck
        ["to-js", arg] -> runWithOne bookPath arg cliToJS
        ["show", arg]  -> runWithOne bookPath arg cliShow
        ["deps", arg]  -> runWithOne bookPath arg cliDeps
        ["rdeps", arg] -> runWithOne bookPath arg cliRDeps
        _              -> printHelp
      case result of
        Left err -> do
          putStrLn err
          exitWith (ExitFailure 1)
        Right _ -> do
          exitWith ExitSuccess

printHelp :: IO (Either String ())
printHelp = do
  putStrLn "Kind usage:"
  putStrLn "  kind check             # Checks all .kind files in the current directory and subdirectories"
  putStrLn "  kind check <name|path> # Type-checks all definitions in the specified file"
  putStrLn "  kind run   <name|path> # Normalizes the specified definition"
  putStrLn "  kind show  <name|path> # Stringifies the specified definition"
  putStrLn "  kind to-js <name|path> # Compiles the specified definition to JavaScript"
  putStrLn "  kind deps  <name|path> # Shows immediate dependencies of the specified definition"
  putStrLn "  kind rdeps <name|path> # Shows all dependencies of the specified definition recursively"
  putStrLn "  kind help              # Shows this help message"
  return $ Right ()

-- CLI Commands
-- ------------

-- Normalizes the target definition
cliNormal :: Command
cliNormal bookPath (book, _, _) defName defPath =
  case M.lookup "main" book of
    Just term -> do
      result <- showInfo book IM.empty (Print term 0)
      putStrLn result
      return $ Right ()
    Nothing -> do
      return $ Left $ "Error: Definition '" ++ defName ++ "' not found."

-- Checks all definitions in the target file
cliCheck :: Command
cliCheck bookPath (book, defs, _) defName defPath = do
  case M.lookup defPath defs of
    Just fileDefNames -> do
      results <- forM fileDefNames $ \fileDefName -> do
        case M.lookup fileDefName book of
          Just term -> do
            case envRun (doCheck term) book of
              Done state _ -> do
                cliPrintLogs state
                cliPrintWarn term state
                putStrLn $ "\x1b[32m✓ " ++ fileDefName ++ "\x1b[0m"
                return $ Right ()
              Fail state -> do
                cliPrintLogs state
                cliPrintWarn term state
                putStrLn $ "\x1b[31m✗ " ++ fileDefName ++ "\x1b[0m"
                return $ Left $ "Error."
          Nothing -> return $ Left $ "Definition not found: " ++ fileDefName
      putStrLn ""
      return $ sequence_ results
    Nothing -> do
      return $ Left $ "No definitions found in file: " ++ defPath

-- Compiles the whole book to JS
cliToJS :: Command
cliToJS bookPath (book, _, _) _ _ = do
  putStrLn $ compileJS book
  return $ Right ()

-- Shows a definition
cliShow :: Command
cliShow bookPath (book, _, _) defName _ = 
  case M.lookup defName book of
    Just term -> do
      putStrLn $ showTerm term
      return $ Right ()
    Nothing -> do
      return $ Left $ "Error: Definition '" ++ defName ++ "' not found."

-- Shows immediate dependencies of a definition
cliDeps :: Command
cliDeps bookPath (book, _, _) defName _ = 
  case M.lookup defName book of
    Just term -> do
      forM_ (filter (/= defName) $ nub $ getDeps term) $ \dep -> putStrLn dep
      return $ Right ()
    Nothing -> do
      return $ Left $ "Error: Definition '" ++ defName ++ "' not found."

-- Shows all dependencies of a definition recursively
cliRDeps :: Command
cliRDeps bookPath (book, _, _) defName _ = do
  let deps = S.toList $ S.delete defName $ getAllDeps book defName
  forM_ deps $ \dep -> putStrLn dep
  return $ Right ()

-- CLI Runners
-- -----------

-- Runs a command on a single file
runWithOne :: FilePath -> String -> Command -> IO (Either String ())
runWithOne bookPath arg action = do
  let defName = getDefName bookPath arg
  let defPath = getDefPath bookPath defName
  cliCtx <- loadName bookPath M.empty defName
  action bookPath cliCtx defName defPath

-- Runs a command on all files
runWithAll :: FilePath -> Command -> IO (Either String ())
runWithAll bookPath action = do
  files <- findKindFiles bookPath
  results <- forM files $ \file -> do
    putStrLn $ "\x1b[1m\x1b[4m[" ++ file ++ "]\x1b[0m"
    runWithOne bookPath file action
  return $ sequence_ results

-- Loader
-- ------

-- Loads a name and all its dependencies recursively
loadName :: FilePath -> Book -> String -> IO FileCtx
loadName bookPath book name = do
  if M.member name book
    then do
      return (book, M.empty, M.empty)
    else do
      let dirPath = bookPath </> name
      isDir <- doesDirectoryExist dirPath
      if isDir
        then loadFile bookPath book (dirPath </> takeFileName name ++ ".kind")
        else loadFile bookPath book (bookPath </> name ++ ".kind")

-- Loads a file and all its dependencies recursivelly
loadFile :: FilePath -> Book -> FilePath -> IO FileCtx
loadFile bookPath book filePath = do
  fileExists <- doesFileExist filePath
  if not fileExists
    then do
      return (book, M.empty, M.empty)
    else do
      code  <- readFile filePath
      book0 <- doParseBook filePath code
      let book1 = M.union book book0
      let defs  = M.keys book0
      let deps  = concatMap (getDeps . snd) (M.toList book0)
      let defs' = M.singleton filePath defs
      let deps' = M.singleton filePath deps
      foldM (\ (depBook, depDefs, depDeps) dep -> do
          (depBook', depDefs', depDeps') <- loadName bookPath depBook dep
          return ( depBook' , M.union depDefs depDefs' , M.union depDeps depDeps')
        ) (book1, defs', deps') deps

-- Utils
-- -----

-- Finds the directory named "monobook"
findBookDir :: FilePath -> IO (Maybe FilePath)
findBookDir dir = do
  let kindBookDir = dir </> "kindbook"
  foundKindBook <- doesDirectoryExist kindBookDir
  if foundKindBook
    then return $ Just kindBookDir
    else if takeDirectory dir == dir
      then return Nothing
      else findBookDir (takeDirectory dir)

-- Finds all Kind files in this directory tree
findKindFiles :: FilePath -> IO [FilePath]
findKindFiles dir = do
  contents <- getDirectoryContents dir
  let properNames = filter (`notElem` [".", ".."]) contents
  paths <- forM properNames $ \name -> do
    let path = dir </> name
    isDirectory <- doesDirectoryExist path
    if isDirectory
      then findKindFiles path
      else return [path | ".kind" `isSuffixOf` path]
  return (concat paths)

-- Loads a file into a string
readSource :: FilePath -> IO String
readSource file = do
  result <- try (readFile file) :: IO (Either IOError String)
  case result of
    Right x -> return x
    Left er -> return $ "Could not read source file: " ++ file

-- Extracts the definition name from a file path or name
getDefName :: FilePath -> String -> String
getDefName bookPath = dropBookPath . dropExtension where
  dropExtension path
    | isExtensionOf "kind" path = System.FilePath.dropExtension path
    | otherwise                 = path
  dropBookPath path = maybe path id (stripPrefix (bookPath++"/") path)

-- Gets the full path for a definition
getDefPath :: FilePath -> String -> FilePath
getDefPath bookPath name = bookPath </> name ++ ".kind"

-- Stringification
-- ---------------

showInfo :: Book -> Fill -> Info -> IO String
showInfo book fill info = case info of
  Found nam typ ctx dep ->
    let nam' = concat ["?", nam]
        typ' = showTermGo True (normal book fill 0 typ dep) dep
        ctx' = showContext book fill ctx dep
    in return $ concat ["\x1b[1mGOAL\x1b[0m ", nam', " : ", typ', "\n", ctx']
  Error src exp det bad dep -> do
    let exp' = concat ["- expected : \x1b[32m", showTermGo True (normal book fill 0 exp dep) dep, "\x1b[0m"]
        det' = concat ["- detected : \x1b[31m", showTermGo True (normal book fill 0 det dep) dep, "\x1b[0m"]
        bad' = concat ["- origin   : \x1b[2m", showTermGo True (normal book fill 0 bad dep) dep, "\x1b[0m"]
    (file, text) <- case src of
      Just (Cod (Loc fileName iniLine iniCol) (Loc _ endLine endCol)) -> do
        canonPath <- canonicalizePath fileName
        content   <- readSource canonPath
        let highlighted = highlightError (iniLine, iniCol) (endLine, endCol) content
        return (canonPath, unlines $ take 8 $ lines highlighted)
      Nothing -> return ("unknown_file", "Could not read source file.\n")
    let src' = concat ["\x1b[4m", file, "\x1b[0m\n", text]
    return $ concat ["\x1b[1mERROR:\x1b[0m\n", exp', "\n", det', "\n", bad', "\n", src']
  Solve nam val dep ->
    return $ concat ["SOLVE: _", show nam, " = ", showTermGo True val dep]
  Vague nam ->
    return $ concat ["VAGUE: _", nam]
  Print val dep ->
    return $ showTermGo True (normal book fill 2 val dep) dep

showContext :: Book -> Fill -> [Term] -> Int -> String
showContext book fill ctx dep = unlines $ map (\term -> "- " ++ showContextAnn book fill term dep) ctx

showContextAnn :: Book -> Fill -> Term -> Int -> String
showContextAnn book fill (Ann chk val typ) dep = concat [showTermGo True (normal book fill 0 val dep) dep, " : ", showTermGo True (normal book fill 0 typ dep) dep]
showContextAnn book fill (Src _ val)       dep = showContextAnn book fill val dep
showContextAnn book fill term              dep = showTermGo True (normal book fill 0 term dep) dep

-- Prints logs from the type-checker
cliPrintLogs :: State -> IO ()
cliPrintLogs (State book fill susp logs) = do
  forM_ logs $ \log -> do
    result <- showInfo book fill log
    putStr result

-- Prints a warning if there are unsolved metas
cliPrintWarn :: Term -> State -> IO ()
cliPrintWarn term (State _ fill _ _) = do
  let metaCount = countMetas term
  let fillCount = IM.size fill
  if (metaCount > fillCount) then do
    putStrLn $ "WARNING: " ++ show (metaCount - fillCount) ++ " unsolved metas."
  else
    return ()


================================================
FILE: src/Kind/Check.hs
================================================
-- //./Type.hs//

module Kind.Check where

import Kind.Env
import Kind.Equal
import Kind.Reduce
import Kind.Show
import Kind.Type
import Kind.Util

import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M

import Control.Monad (forM, forM_, unless, when)
import Debug.Trace

-- Type-Checking
-- -------------

-- Modes:
-- - sus=True  : suspended checks on / better unification / wont return annotated term 
-- - sus=False : suspended checks off / worse unification / will return annotated term

infer :: Bool -> Maybe Cod -> Term -> Int -> Env Term
infer sus src term dep = debug ("infer:" ++ (if sus then "* " else " ") ++ showTermGo False term dep) $ go term where

  go (All nam inp bod) = do
    inpA <- checkLater sus src inp Set dep
    bodA <- checkLater sus src (bod (Ann False (Var nam dep) inp)) Set (dep + 1)
    return $ Ann False (All nam inpA (\x -> bodA)) Set

  go (App fun arg) = do
    funA <- infer sus src fun dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType funA) of
      (All inpNam inpTyp inpBod) -> do
        argA <- checkLater sus src arg inpTyp dep
        return $ Ann False (App funA argA) (inpBod arg)
      otherwise -> do
        envLog (Error src (Ref "function") (getType funA) (App fun arg) dep)
        envFail

  go (Ann True val typ) = do
    check sus src val typ dep

  go (Ann False val typ) = do
    return $ Ann False val typ

  go (Slf nam typ bod) = do
    typA <- checkLater sus src typ Set dep
    bodA <- checkLater sus src (bod (Ann False (Var nam dep) typ)) Set (dep + 1)
    return $ Ann False (Slf nam typA (\x -> bodA)) Set

  go (Ins val) = do
    valA <- infer sus src val dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType valA) of
      (Slf slfNam slfTyp slfBod) -> do
        return $ Ann False (Ins valA) (slfBod (Ins valA))
      otherwise -> do
        envLog (Error src (Ref "Self") (getType valA) (Ins val) dep)
        envFail

  go (Ref nam) = do
    book <- envGetBook
    case M.lookup nam book of
      Just val -> do
        valA <- infer sus src val dep
        return $ Ann False (Ref nam) (getType valA)
      Nothing -> do
        envLog (Error src (Ref "expression") (Ref "undefined") (Ref nam) dep)
        envFail

  go Set = do
    return $ Ann False Set Set

  go U64 = do
    return $ Ann False U64 Set

  go F64 = do
    return $ Ann False F64 Set

  go (Num num) = do
    return $ Ann False (Num num) U64

  go (Flt num) = do
    return $ Ann False (Flt num) F64


  go (Op2 opr fst snd) = do
    fstT <- infer sus src fst dep
    sndT <- infer sus src snd dep
 
    let validTypes = [F64, U64]
    isValidType <- checkValidType (getType fstT) validTypes dep

    if not isValidType then do
      envLog (Error src (Ref "Valid numeric type") (getType fstT) (Op2 opr fst snd) dep)
      envFail
    else do
      typesEqual <- equal (getType fstT) (getType sndT) dep
      if not typesEqual then do
        envLog (Error src (getType fstT) (getType sndT) (Op2 opr fst snd) dep)
        envFail
      else do
        book <- envGetBook
        fill <- envGetFill
        let reducedFst = reduce book fill 1 (getType fstT)
        let returnType = getOpReturnType opr reducedFst
        return $ Ann False (Op2 opr fstT sndT) returnType
  
  go (Swi zer suc) = do
    envLog (Error src (Ref "annotation") (Ref "switch") (Swi zer suc) dep)
    envFail

  go (Map typ) = do
    typA <- checkLater sus src typ Set dep
    return $ Ann False (Map typA) Set

  go (KVs kvs dft) = do
    dftA <- infer sus src dft dep
    kvsA <- forM (IM.toList kvs) $ \(key, val ) -> do
      valA <- check sus src val (getType dftA) dep
      return (key, valA)
    return $ Ann False (KVs (IM.fromList kvsA) dftA) (Map (getType dftA))

  go (Get got nam map key bod) = do
    mapA <- infer sus src map dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType mapA) of
      (Map typ) -> do
        let got_ann = Ann False (Var got dep) typ
        let nam_ann = Ann False (Var nam dep) (Map typ)
        keyA <- check sus src key U64 dep
        bodA <- infer sus src (bod got_ann nam_ann) dep
        return $ Ann False (Get got nam mapA keyA (\g m -> bodA)) (getType bodA)
      otherwise -> do
        envLog (Error src (Ref "Map") (getType mapA) (Get got nam map key bod) dep)
        envFail

  go (Put got nam map key val bod) = do
    mapA <- infer sus src map dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType mapA) of
      (Map typ) -> do
        valA <- check sus src val typ dep
        let got_ann = Ann False (Var got dep) typ
        let nam_ann = Ann False (Var nam dep) (Map typ)
        keyA <- check sus src key U64 dep
        bodA <- infer sus src (bod got_ann nam_ann) dep
        return $ Ann False (Put got nam mapA keyA valA (\g m -> bodA)) (getType bodA)
      otherwise -> do
        envLog (Error src (Ref "Map") (getType mapA) (Put got nam map key val bod) dep)
        envFail

  go (Let nam val bod) = do
    valA <- infer sus src val dep
    bodA <- infer sus src (bod (Ann False (Var nam dep) (getType valA))) dep
    return $ Ann False (Let nam valA (\x -> bodA)) (getType bodA)

  go (Use nam val bod) = do
    infer sus src (bod val) dep

  -- TODO: annotate inside ADT for completion (not needed)
  go (ADT scp cts typ) = do
    ctsA <- forM cts $ \ (Ctr cnam tele) -> do
      teleA <- checkTele sus src tele Set dep
      return $ Ctr cnam teleA
    return $ Ann False (ADT scp ctsA typ) Set

  go (Con nam arg) = do
    envLog (Error src (Ref "annotation") (Ref "constructor") (Con nam arg) dep)
    envFail

  go (Mat cse) = do
    envLog (Error src (Ref "annotation") (Ref "match") (Mat cse) dep)
    envFail

  go (Lam nam bod) = do
    envLog (Error src (Ref "annotation") (Ref "lambda") (Lam nam bod) dep)
    envFail

  go (Hol nam ctx) = do
    envLog (Error src (Ref "annotation") (Ref "hole") (Hol nam ctx) dep)
    envFail

  go (Met uid spn) = do
    envLog (Error src (Ref "annotation") (Ref "meta") (Met uid spn) dep)
    envFail

  go (Log msg nxt) = do
    msgA <- infer sus src msg dep
    nxtA <- infer sus src nxt dep
    return $ Ann False (Log msgA nxtA) (getType nxtA)

  go (Var nam idx) = do
    envLog (Error src (Ref "annotation") (Ref "variable") (Var nam idx) dep)
    envFail

  go (Src src val) = do
    infer sus (Just src) val dep

  go tm@(Txt txt) = do
    return $ Ann False tm (Ref "String")
    -- book <- envGetBook
    -- fill <- envGetFill
    -- go (reduce book fill 2 tm)

  go tm@(Nat val) = do
    book <- envGetBook
    fill <- envGetFill
    go (reduce book fill 2 tm)

  go tm@(Lst lst) = do
    book <- envGetBook
    fill <- envGetFill
    go (reduce book fill 2 tm)

check :: Bool -> Maybe Cod -> Term -> Term -> Int -> Env Term
check sus src term typx dep = debug ("check:" ++ (if sus then "* " else " ") ++ showTermGo False term dep ++ "\n    :: " ++ showTermGo True typx dep) $ go term where

  go (App (Src _ val) arg) =
    go (App val arg)

  go (App (Mat cse) arg) = do
    argA <- infer sus src arg dep
    infer sus src (App (Ann True (Mat cse) (All "x" (getType argA) (\x -> replace arg x typx dep))) arg) dep

  go (App (Swi zer suc) arg) = do
    argA <- infer sus src arg dep
    infer sus src (App (Ann True (Swi zer suc) (All "x" (getType argA) (\x -> replace arg x typx dep))) arg) dep

  go (Lam nam bod) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      (All typNam typInp typBod) -> do
        let ann = Ann False (Var nam dep) typInp
        bodA <- check sus src (bod ann) (typBod ann) (dep + 1)
        return $ Ann False (Lam nam (\x -> bodA)) typx
      otherwise -> do
        infer sus src (Lam nam bod) dep

  go (Ins val) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      Slf typNam typTyp typBod -> do
        valA <- check sus src val (typBod (Ins val)) dep
        return $ Ann False (Ins valA) typx
      _ -> infer sus src (Ins val) dep

  go val@(Con nam arg) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      (ADT adtScp adtCts adtTyp) -> do
        case lookup nam (map (\(Ctr cNam cTel) -> (cNam, cTel)) adtCts) of
          Just cTel -> do
            argA <- checkConstructor src arg cTel dep
            return $ Ann False (Con nam argA) typx
          Nothing -> do
            envLog (Error src (Hol ("constructor_not_found:"++nam) []) (Hol "unknown_type" []) (Con nam arg) dep)
            envFail
      otherwise -> infer sus src (Con nam arg) dep
    where
      checkConstructor :: Maybe Cod -> [(Maybe String, Term)] -> Tele -> Int -> Env [(Maybe String, Term)]
      checkConstructor src [] (TRet ret) dep = do
        cmp src val ret typx dep
        return []
      checkConstructor src ((field, arg):args) (TExt nam inp bod) dep =
        case field of
          Just field -> if field /= nam
            then do
              envLog (Error src (Hol ("expected:" ++ nam) []) (Hol ("detected:" ++ field) []) (Hol "field_mismatch" []) dep)
              envFail
            else do
              argA  <- check sus src arg inp dep
              argsA <- checkConstructor src args (bod arg) (dep + 1)
              return $ (Just field, argA) : argsA
          Nothing -> do
            argA  <- check sus src arg inp dep
            argsA <- checkConstructor src args (bod arg) (dep + 1)
            return $ (Nothing, argA) : argsA
      checkConstructor src _ _ dep = do
        envLog (Error src (Hol "arity_mismatch" []) (Hol "unknown_type" []) (Hol "constructor" []) dep)
        envFail

  go (Mat cse) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      (All typNam typInp typBod) -> do
        case reduce book fill 2 typInp of
          (ADT adtScp adtCts adtTyp) -> do
            -- Checks if all cases are well-typed
            let adtCtsMap = M.fromList (map (\ (Ctr cNam cTel) -> (cNam, cTel)) adtCts)
            let coveredCases = M.fromList cse
            cseA <- forM cse $ \ (cNam, cBod) -> do
              if cNam == "_" then do
                if null (adtCtsMap `M.difference` coveredCases) then do
                  checkUnreachable Nothing cNam cBod dep
                else do
                  cBodA <- check sus src cBod (All "" typInp typBod) dep
                  return (cNam, cBodA)
              else case M.lookup cNam adtCtsMap of
                Just cTel -> do
                  let a_r = teleToTerms cTel dep
                  let eqs = zip (getDatIndices (reduce book fill 2 typInp)) (getDatIndices (reduce book fill 2 (snd a_r)))
                  let rt0 = teleToType cTel (typBod (Ann False (Con cNam (fst a_r)) typInp)) dep
                  let rt1 = foldl' (\ ty (a,b) -> replace a b ty dep) rt0 eqs
                  if any (\(a,b) -> incompatible a b dep) eqs then
                    checkUnreachable Nothing cNam cBod dep
                  else do
                    cBodA <- check sus src cBod rt1 dep
                    return (cNam, cBodA)
                Nothing -> do
                  envLog (Error src (Hol ("constructor_not_found:"++cNam) []) (Hol "unknown_type" []) (Mat cse) dep)
                  envFail
            -- Check if all constructors are covered
            forM_ adtCts $ \ (Ctr cNam _) ->
              unless (M.member cNam coveredCases || M.member "_" coveredCases) $ do
                envLog (Error src (Hol ("missing_case:" ++ cNam) []) (Hol "incomplete_match" []) (Mat cse) dep)
                envFail
            return $ Ann False (Mat cseA) typx
          otherwise -> infer sus src (Mat cse) dep
      otherwise -> infer sus src (Mat cse) dep

  go (Swi zer suc) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      (All typNam typInp typBod) -> do
        case reduce book fill 2 typInp of
          U64 -> do
            -- Check zero case
            let zerAnn = Ann False (Num 0) U64
            zerA <- check sus src zer (typBod zerAnn) dep
            -- Check successor case
            let sucAnn = Ann False (Var "n" dep) U64
            let sucTyp = All "n" U64 (\x -> typBod (Op2 ADD (Num 1) x))
            sucA <- check sus src suc sucTyp dep
            return $ Ann False (Swi zerA sucA) typx
          otherwise -> infer sus src (Swi zer suc) dep
      otherwise -> infer sus src (Swi zer suc) dep

  go (KVs kvs dft) = do
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 typx of
      (Map typ) -> do
        dftA <- check sus src dft typ dep
        kvsA <- forM (IM.toList kvs) $ \(key, val) -> do
          valA <- check sus src val typ dep
          return (key, valA)
        return $ Ann False (KVs (IM.fromList kvsA) dftA) typx
      otherwise -> infer sus src (KVs kvs dft) dep

  go (Get got nam map key bod) = do
    mapA <- infer sus src map dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType mapA) of
      (Map typ) -> do
        let got_ann = Ann False (Var got dep) typ
        let nam_ann = Ann False (Var nam dep) (Map typ)
        keyA <- check sus src key U64 dep
        bodA <- check sus src (bod got_ann nam_ann) typx dep
        return $ Ann False (Get got nam mapA keyA (\g m -> bodA)) typx
      otherwise -> do
        envLog (Error src (Ref "Map") (getType mapA) (Get got nam map key bod) dep)
        envFail

  go (Put got nam map key val bod) = do
    mapA <- infer sus src map dep
    book <- envGetBook
    fill <- envGetFill
    case reduce book fill 2 (getType mapA) of
      (Map typ) -> do
        valA <- check sus src val typ dep
        let got_ann = Ann False (Var got dep) typ
        let nam_ann = Ann False (Var nam dep) (Map typ)
        keyA <- check sus src key U64 dep
        bodA <- check sus src (bod got_ann nam_ann) typx dep
        return $ Ann False (Put got nam mapA keyA valA (\g m -> bodA)) typx
      otherwise -> do
        envLog (Error src (Ref "Map") (getType mapA) (Put got nam map key val bod) dep)
        envFail

  go (Let nam val bod) = do
    valA <- infer sus src val dep
    bodA <- check sus src (bod (Ann False (Var nam dep) (getType valA))) typx dep
    return $ Ann False (Let nam valA (\x -> bodA)) typx

  go (Use nam val bod) = do
    check sus src (bod val) typx dep

  go (Hol nam ctx) = do
    envLog (Found nam typx ctx dep)
    return $ Ann False (Hol nam ctx) typx

  go (Met uid spn) = do
    return $ Ann False (Met uid spn) typx

  go (Log msg nxt) = do
    msgA <- infer sus src msg dep
    nxtA <- check sus src nxt typx dep
    return $ Ann False (Log msgA nxtA) typx

  go tm@(Txt txt) = do
    return $ Ann False tm (Ref "String")
    -- book <- envGetBook
    -- fill <- envGetFill
    -- go (reduce book fill 2 tm)

  go tm@(Nat val) = do
    book <- envGetBook
    fill <- envGetFill
    go (reduce book fill 2 tm)

  go tm@(Lst lst) = do
    book <- envGetBook
    fill <- envGetFill
    go (reduce book fill 2 tm)

  go (Ann True val typ) = do
    cmp src val typ typx dep
    check sus src val typ dep

  go (Ann False val typ) = do
    cmp src val typ typx dep -- FIXME: should this be here?
    return $ Ann False val typ

  go (Src src val) = do
    check sus (Just src) val typx dep

  go term = do
    termA <- infer sus src term dep
    cmp src term typx (getType termA) dep
    return termA

  cmp src term expected detected dep = do
    equal <- equal expected detected dep
    if equal then do
      susp <- envTakeSusp
      forM_ susp $ \ (Check src val typ dep) -> do
        check sus src val typ dep
      return ()
    else do
      envLog (Error src expected detected term dep)
      envFail

checkTele :: Bool -> Maybe Cod -> Tele -> Term -> Int -> Env Tele
checkTele sus src tele typ dep = case tele of
  TRet term -> do
    termA <- check sus src term typ dep
    return $ TRet termA
  TExt nam inp bod -> do
    inpA <- check sus src inp Set dep
    bodA <- checkTele sus src (bod (Ann False (Var nam dep) inp)) typ (dep + 1)
    return $ TExt nam inpA (\x -> bodA)

checkUnreachable :: Maybe Cod -> String -> Term -> Int -> Env (String, Term)
checkUnreachable src cNam term dep = go src cNam term dep where
  go src cNam (Lam nam bod)     dep = go src cNam (bod (Con "void" [])) (dep+1)
  go src cNam (Let nam val bod) dep = go src cNam (bod (Con "void" [])) (dep+1)
  go src cNam (Use nam val bod) dep = go src cNam (bod (Con "void" [])) (dep+1)
  go _   cNam (Src src val)     dep = go (Just src) cNam val dep
  go src cNam (Hol nam ctx)     dep = envLog (Found nam (Hol "unreachable" []) ctx dep) >> go src cNam Set dep
  go src cNam term              dep = return (cNam, Ann False Set U64)

checkLater :: Bool -> Maybe Cod -> Term -> Term -> Int -> Env Term
checkLater False src term typx dep = check False src term typx dep
checkLater True  src term typx dep = envSusp (Check src term typx dep) >> return (Met 0 [])

doCheckMode :: Bool -> Term -> Env Term
doCheckMode sus (Ann _ val typ) = do
  check sus Nothing typ Set 0
  check sus Nothing val typ 0
doCheckMode sus (Src _ val) = do
  doCheckMode sus val
doCheckMode sus (Ref nam) = do
  book <- envGetBook
  case M.lookup nam book of
    Just val -> doCheckMode sus val
    Nothing  -> envLog (Error Nothing (Ref "expression") (Ref "undefined") (Ref nam) 0) >> envFail
doCheckMode sus term = do
  infer True Nothing term 0

doCheck :: Term -> Env Term
doCheck = doCheckMode True

doAnnotate :: Term -> Env (Term, Fill)
doAnnotate term = do
  doCheckMode True term
  term <- doCheckMode False term
  fill <- envGetFill
  return (bind term [], fill)


================================================
FILE: src/Kind/CompileJS.hs
================================================
-- Type.hs:
-- //./Type.hs//

-- FIXME: currently, the Map type will compile to a mutable map in JS, which
-- means we assume it is used linearly (no cloning). To improve this, we can add
-- a shallow-cloning operation for cloned maps, or use an immutable map. Adding
-- linearity checks to Kind would let us pick the best representation.

{-# LANGUAGE ViewPatterns #-}

module Kind.CompileJS where

import Kind.Check
import Kind.Env
import Kind.Equal
import Kind.Reduce
import Kind.Show
import Kind.Type
import Kind.Util

import Control.Monad (forM)
import Data.List (intercalate, isSuffixOf, elem, isInfixOf, isPrefixOf)
import Data.Maybe (fromJust, isJust)
import Data.Word
import qualified Control.Monad.State.Lazy as ST
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Set as S

import Debug.Trace

import Prelude hiding (EQ, LT, GT)

-- Type
-- ----

-- Compilation Targets
data Target = C | JS deriving (Eq, Show)

-- Compilable Term
data CT
  = CNul
  | CSet
  | CU64
  | CF64
  | CADT [(String,[(String,CT)])]
  | CMap CT
  | CAll (String,CT) (CT -> CT)
  | CLam (String,CT) (CT -> CT)
  | CApp CT CT
  | CCon String [(String, CT)]
  | CMat CT [(String, [(String,CT)], CT)]
  | CRef String
  | CHol String
  | CLet (String,CT) CT (CT -> CT)
  | CNum Word64
  | CFlt Double
  | COp2 CT Oper CT CT
  | CSwi CT CT CT
  | CKVs (IM.IntMap CT) CT
  | CGet String String CT CT (CT -> CT -> CT)
  | CPut String String CT CT CT (CT -> CT -> CT)
  | CLog CT CT
  | CVar String Int
  | CTxt String
  | CLst [CT]
  | CNat Integer

type CTBook = M.Map String CT

-- Term to CT
-- ----------

-- Converts a Term into a Compilable Term
-- Uses type information to:
-- - Ensure constructor fields are present
-- - Annotate Mat cases with the field names
termToCT :: Book -> Fill -> Term -> Maybe Term -> Int -> CT
termToCT book fill term typx dep = bindCT (t2ct term typx dep) [] where

  t2ct term typx dep = 
    -- trace ("t2ct: " ++ showTerm term ++ "\ntype: " ++ maybe "*" showTerm typx ++ "\ndep: " ++ show dep) $
    go term where

    go (All nam inp bod) =
      let inp' = t2ct inp Nothing dep
          bod' = \x -> t2ct (bod (Var nam dep)) Nothing (dep+1)
      in CAll (nam,inp') bod'
    go (Lam nam bod) =
      case typx of
        Just typx -> case (reduce book fill 2 typx) of
          (All _ inp _) ->
            let inp' = t2ct inp Nothing dep
                bod' = \x -> t2ct (bod (Var nam dep)) Nothing (dep+1)
            in CLam (nam,inp') bod'
          other -> error "err"
        Nothing -> error "err"
    go (App fun arg) =
      let fun' = t2ct fun Nothing dep
          arg' = t2ct arg Nothing dep
      in CApp fun' arg'
    go (Ann _ val typ) =
      t2ct val (Just typ) dep
    go (Slf _ _ _) =
      CNul
    go (Ins val) =
      t2ct val typx dep
    go (ADT scp cts typ) =
      let cts' = map (\ (Ctr nam tele) -> (nam, map (\ (fn,ft) -> (fn, go ft)) (getTeleFields tele dep []))) cts
      in CADT cts'
    go (Con nam arg) =
      case typx of
        Just typx -> case lookup nam (getADTCts (reduce book fill 2 typx)) of
          Just (Ctr _ tele) ->
            let fNames = getTeleNames tele dep []
                fields = map (\ (f,t) -> (f, t2ct t Nothing dep)) $ zip fNames (map snd arg)
            in CCon nam fields
          Nothing -> error $ "constructor-not-found:" ++ nam
        Nothing -> error $ "untyped-constructor"
    go (Mat cse) =
      case typx of
        Just typx -> case reduce book fill 2 typx of
          (All _ adt _) ->
            let adtV = reduce book fill 2 adt
                cts  = getADTCts adtV
                adt' = t2ct adt Nothing dep
                cses = map (\ (cnam, cbod) ->
                  if cnam == "_" then
                    (cnam, [("_",adt')], t2ct cbod Nothing dep)
                  else case lookup cnam cts of
                    Just (Ctr _ tele) ->
                      let fInps  = getTeleFields tele dep []
                          fInps' = map (\ (nm,ty) -> (nm, t2ct ty Nothing dep)) fInps
                      in (cnam, fInps', t2ct cbod Nothing dep)
                    Nothing -> error $ "constructor-not-found:" ++ cnam) cse
            in CLam ("__" ++ show dep, adt') $ \x -> CMat x cses
          otherwise -> error "match-without-type"
        Nothing -> error "err"
    go (Swi zer suc) =
      let zer' = t2ct zer Nothing dep
          suc' = t2ct suc Nothing dep
      in CLam ("__" ++ show dep, CU64) $ \x -> CSwi x zer' suc'
    go (Map typ) =
      let typ' = t2ct typ Nothing dep
      in CMap typ'
    go (KVs kvs def) =
      let kvs' = IM.map (\v -> t2ct v Nothing dep) kvs
          def' = t2ct def Nothing dep
      in CKVs kvs' def'
    go (Get got nam map key bod) =
      let map' = t2ct map Nothing dep
          key' = t2ct key Nothing dep
          bod' = \x y -> t2ct (bod (Var got dep) (Var nam dep)) Nothing (dep+2)
      in CGet got nam map' key' bod'
    go (Put got nam map key val bod) =
      let map' = t2ct map Nothing dep
          key' = t2ct key Nothing dep
          val' = t2ct val Nothing dep
          bod' = \x y -> t2ct (bod (Var got dep) (Var nam dep)) Nothing (dep+2)
      in CPut got nam map' key' val' bod'
    go (All _ _ _) =
      CNul
    go (Ref nam) =
      CRef nam
    go (Let nam val bod) =
      -- FIXME: add type
      let val' = t2ct val Nothing dep
          bod' = \x -> t2ct (bod (Var nam dep)) Nothing (dep+1)
      in CLet (nam,CNul) val' bod'
    go (Use nam val bod) =
      t2ct (bod val) typx dep
    go Set =
      CSet
    go U64 =
      CU64
    go F64 =
      CF64
    go (Num val) =
      CNum val
    go (Flt val) =
      CFlt val
    go (Op2 opr fst snd) = case typx of
      Nothing -> error "Type information required for binary operation"
      Just typ -> 
          let fst' = t2ct fst Nothing dep
              snd' = t2ct snd Nothing dep
              typ' = t2ct typ Nothing dep
          in COp2 typ' opr fst' snd'
    go (Txt txt) =
      CTxt txt
    go (Lst lst) =
      CLst (map (\x -> t2ct x Nothing dep) lst)
    go (Nat val) =
      CNat val
    go (Hol nam _) =
      CHol nam
    go (Met _ _) =
      CNul
    go (Log msg nxt) =
      let msg' = t2ct msg Nothing dep
          nxt' = t2ct nxt Nothing dep
      in CLog msg' nxt'
    go (Var nam idx) =
      CVar nam idx
    go (Src _ val) =
      t2ct val typx dep

-- CT Transformations
-- ------------------

-- Removes unreachable cases
removeUnreachables :: CT -> CT
removeUnreachables ct = go ct where
  go CNul =
    CNul
  go CSet =
    CSet
  go CU64 =
    CU64
  go CF64 =
    CF64
  go (CADT cts) =
    let cts' = map (\ (n,fs) -> (n, map (\ (fn,ft) -> (fn, go ft)) fs)) cts
    in CADT cts'
  go (CMap typ) =
    let typ' = go typ
    in CMap typ'
  go (CMat val cse) =
    let val' = go val
        cse' = map (\ (n,f,t) -> (n, map (\ (fn,ft) -> (fn, go ft)) f, go t)) cse
        cseF = filter (\ (_,_,t) -> not (isNul t)) cse'
    in CMat val' cseF
  go (CAll (nam,inp) bod) =
    let inp' = go inp
        bod' = \x -> go (bod x)
    in CAll (nam,inp') bod'
  go (CLam (nam,inp) bod) =
    let inp' = go inp
        bod' = \x -> go (bod x)
    in CLam (nam,inp') bod'
  go (CApp fun arg) =
    let fun' = go fun
        arg' = go arg
    in CApp fun' arg'
  go (CCon nam fields) =
    let fields' = map (\ (f,t) -> (f, go t)) fields
    in CCon nam fields'
  go (CRef nam) = CRef nam
  go (CHol nam) = CHol nam
  go (CLet (nam,typ) val bod) =
    let typ' = go typ
        val' = go val
        bod' = \x -> go (bod x)
    in CLet (nam,typ') val' bod'
  go (CNum val) =
    CNum val
  go (CFlt val) =
    CFlt val
  go (COp2 typ opr fst snd) =
    let fst' = go fst
        snd' = go snd
        typ' = go typ
    in COp2 typ' opr fst' snd'
  go (CSwi val zer suc) =
    let val' = go val
        zer' = go zer
        suc' = go suc
    in CSwi val' zer' suc'
  go (CKVs kvs def) =
    let kvs' = IM.map go kvs
        def' = go def
    in CKVs kvs' def'
  go (CGet got nam map key bod) =
    let map' = go map
        key' = go key
        bod' = \x y -> go (bod x y)
    in CGet got nam map' key' bod'
  go (CPut got nam map key val bod) =
    let map' = go map
        key' = go key
        val' = go val
        bod' = \x y -> go (bod x y)
    in CPut got nam map' key' val' bod'
  go (CLog msg nxt) =
    let msg' = go msg
        nxt' = go nxt
    in CLog msg' nxt'
  go (CVar nam idx) =
    CVar nam idx
  go (CTxt txt) =
    CTxt txt
  go (CLst lst) =
    CLst (map go lst)
  go (CNat val) =
    CNat val

-- Lifts shareable lambdas across branches:
-- - from: λx       match v { #Foo{a b}: λy λz A #Bar: λy λz B ... }
-- -   to: λx λy λz match v { #Foo{a b}:       A #Bar:       B ... }
-- TODO: document why this is (and has to be) terrible
-- NOTE: this loses dependencies, turning foralls into simple arrows
liftLambdas :: CT -> Int -> CT
liftLambdas ct depth = 
  gen (liftInp ct depth [] 0) [] ct depth where

  gen :: [CT] -> [CT] -> CT -> Int -> CT
  gen []         ctx ct dep = liftVal ctx ct dep [] 0
  gen (inp:inps) ctx ct dep = CLam (nam dep, inp) (\x -> gen inps (ctx++[x]) ct (dep+1))

  nam :: Int -> String
  nam d = "_" ++ "$" ++ show d

  var :: [CT] -> Int -> CT
  var ctx d | d < length ctx = ctx !! d
  var ctx d | otherwise      = CNul

  eta :: [(String,CT)] -> CT -> CT
  eta []         ct                   = ct
  eta (fld:flds) (CLam (nam,inp) bod) = CLam (nam,inp) $ \x -> eta flds (bod x)
  eta (fld:flds) ct                   = CLam fld       $ \x -> CApp (eta flds ct) x

  liftVal :: [CT] -> CT -> Int -> [CT] -> Int -> CT
  liftVal ctx ct dep inp skip = go ct dep inp skip where
    go (CLam (nam,inp) bod)     dep inps 0    = liftVal ctx (bod (var ctx (length inps))) (dep+1) (inps++[inp]) 0
    go (CLam (nam,inp) bod)     dep inps skip = CLam (nam,inp) $ \x -> liftVal ctx (bod x) (dep+1) inps (skip-1)
    go (CLet (nam,typ) val bod) dep inps skip = CLet (nam,typ) val $ \x -> liftVal ctx (bod x) (dep+1) inps skip
    go ct@(CMat val cse)     dep inps skip | length cse > 0 =
      let recsV = flip map cse $ \ (_,f,b) -> liftVal ctx (eta f b) dep inps (skip + length f)
          recsI = flip map cse $ \ (_,f,b) -> liftInp     (eta f b) dep inps (skip + length f)
          valid = flip all recsI $ \ a -> length a == length (head recsI)
      in if valid then CMat val (zipWith (\ (n,f,_) b -> (n,f,b)) cse recsV) else ct
    go ct@(CSwi val zer suc) dep inps skip =
      let recZI = liftInp     (eta []           zer) dep inps skip
          recZV = liftVal ctx (eta []           zer) dep inps skip
          recSI = liftInp     (eta [("p",CU64)] suc) dep inps (skip + 1)
          recSV = liftVal ctx (eta [("p",CU64)] suc) dep inps (skip + 1)
          valid = length recZI == length recSI
      in if valid then CSwi val recZV recSV else ct
    go ct dep inps s = ct

  liftInp :: CT -> Int -> [CT] -> Int -> [CT]
  liftInp ct dep inps skip = go ct dep inps skip where
    go (CLam (nam,inp) bod)     dep inps 0    = liftInp (bod CNul) (dep+1) (inps++[inp]) 0
    go (CLam (nam,inp) bod)     dep inps skip = liftInp (bod CNul) (dep+1) inps (skip-1)
    go (CLet (nam,typ) val bod) dep inps skip = liftInp (bod CNul) (dep+1) inps skip
    go (CMat val cse)           dep inps skip | length cse > 0 =
      let recsI = flip map cse $ \ (_,f,b) -> liftInp (eta f b) dep inps (skip + length f)
          valid = flip all recsI $ \ a -> length a == length (head recsI)
      in if valid then head recsI else inps
    go (CSwi val zer suc) dep inps skip =
      let recZI = liftInp (eta []           zer) dep inps skip
          recSI = liftInp (eta [("p",CU64)] suc) dep inps (skip + 1)
          valid = length recZI == length recSI
      in if valid then recZI else inps
    go ct dep inps s = inps

inline :: CTBook -> CT -> CT
inline book ct = nf ct where
  nf :: CT -> CT
  nf ct = go (red book ct) where
    go :: CT -> CT
    go CNul                     = CNul
    go CSet                     = CSet
    go CU64                     = CU64
    go CF64                     = CF64
    go (CADT cts)               = CADT (map (\ (n,fs) -> (n, map (\ (fn,ft) -> (fn, nf ft)) fs)) cts)
    go (CMap typ)               = CMap (nf typ)
    go (CAll (nam,inp) bod)     = CAll (nam, nf inp) (\x -> nf (bod x))
    go (CLam (nam,inp) bod)     = CLam (nam, nf inp) (\x -> nf (bod x))
    go (CApp fun arg)           = CApp (nf fun) (nf arg)
    go (CCon nam fields)        = CCon nam (map (\ (f,t) -> (f, nf t)) fields)
    go (CADT cts)               = CADT (map (\ (n,fs) -> (n, map (\ (fn,ft) -> (fn, nf ft)) fs)) cts)
    go (CMat val cses)          = CMat (nf val) (map (\ (n,f,b) -> (n, map (\ (fn,ft) -> (fn, nf ft)) f, nf b)) cses)
    go (CRef nam)               = CRef nam
    go (CHol nam)               = CHol nam
    go (CLet (nam,typ) val bod) = CLet (nam, nf typ) (nf val) (\x -> nf (bod x))
    go (CNum val)               = CNum val
    go (CFlt val)               = CFlt val
    go (COp2 typ opr fst snd)   = COp2 (nf typ) opr (nf fst) (nf snd)
    go (CSwi val zer suc)       = CSwi (nf val) (nf zer) (nf suc)
    go (CKVs kvs def)           = CKVs (IM.map nf kvs) (nf def)
    go (CGet g n m k b)         = CGet g n (nf m) (nf k) (\x y -> nf (b x y))
    go (CPut g n m k v b)       = CPut g n (nf m) (nf k) (nf v) (\x y -> nf (b x y))
    go (CLog msg nxt)           = CLog (nf msg) (nf nxt)
    go (CVar nam idx)           = CVar nam idx
    go (CTxt txt)               = CTxt txt
    go (CLst lst)               = CLst (map nf lst)
    go (CNat val)               = CNat val

-- CT Evaluation
-- -------------

-- Reduce to WNF
red :: CTBook -> CT -> CT
red book tm = go tm where
  go (CApp fun arg) = app book (red book fun) arg
  go (CRef nam)     = ref book nam    
  go val            = val

-- (let x = y A B)
-- ---------------
-- let x = y (A B)

-- Application
app :: CTBook -> CT -> CT -> CT
app book (CAll (nam,inp) bod)     arg = red book (bod (red book arg))
app book (CLam (nam,inp) bod)     arg = red book (bod (red book arg))
app book (CMat val cse)           arg = CMat val (map (\ (n,f,b) -> (n, f, skp f b (\b -> CApp b arg))) cse)
app book (CLet (nam,typ) val bod) arg = CLet (nam,typ) val (\x -> app book (bod x) arg)
app book fun                      arg = CApp fun arg

-- Maps inside N lambdas
skp :: [(String,CT)] -> CT -> (CT -> CT) -> CT
skp []         ct fn = fn ct
skp (fld:flds) ct fn = CLam fld $ \x -> skp flds (CApp ct x) fn

-- Reference
-- NOTE: this should only inline refs ending with "bind", "bind/go" or "pure".
-- create an aux function called "inl :: String -> Bool" after it
ref :: CTBook -> String -> CT
ref book nam
  | inl nam   = red book (fromJust (M.lookup nam book))
  | otherwise = CRef nam
  where
    inl :: String -> Bool
    inl nam = any (`isSuffixOf` nam)
      [ "/bind"
      , "/bind/go"
      , "/pure"
      -- , "HVM/RTag/eq"
      -- , "HVM/RTerm/get-lab"
      -- , "HVM/RTerm/get-loc"
      -- , "HVM/RTerm/get-tag"
      -- , "HVM/RTerm/new"
      -- , "HVM/alloc-redex"
      -- , "HVM/alloc-rnod"
      -- , "HVM/get"
      -- , "HVM/just"
      -- , "HVM/link"
      -- , "HVM/port"
      -- , "HVM/push-redex"
      -- , "HVM/set"
      -- , "HVM/swap"
      -- , "HVM/take"
      -- , "U64/to-bool"
      , "IO/print"
      , "IO/prompt"
      , "IO/swap"
      , "IO/read"
      , "IO/exec"
      , "IO/args"
      ]

-- JavaScript Codegen
-- ------------------

getArguments :: CT -> ([(String,CT)], CT)
getArguments term = go term 0 where
  go (CLam (nam,inp) bod) dep =
    let (args, body) = go (bod (CVar nam dep)) (dep+1)
    in ((nam,inp):args, body)
  go body dep = ([], body)

arityOf :: CTBook -> String -> Int
arityOf book name = case M.lookup name book of
  Just ct -> length $ fst $ getArguments ct
  Nothing -> 0

isRecCall :: String -> Int -> CT -> [CT] -> Bool
isRecCall fnName arity appFun appArgs =
  case appFun of
    CRef appFunName ->
      let isSameFunc  = appFunName == fnName
          isSameArity = length appArgs == arity
      in isSameFunc && isSameArity
    _ -> False

isSatCall :: CTBook -> CT -> [CT] -> Bool
isSatCall book (CRef funName) appArgs = arityOf book funName == length appArgs
isSatCall book _              _       = False

isEffCall :: CTBook -> CT -> [CT] -> Bool
isEffCall book (CHol name) appArgs = True
isEffCall book name        appArgs = False

-- Converts a function to JavaScript or C
fnToJS :: CTBook -> String -> CT -> ST.State Int String
fnToJS book fnName ct@(getArguments -> (fnArgs, fnBody)) = do
  bodyName <- fresh
  bodyStmt <- ctToJS True bodyName fnBody 0 
  argTypes <- return $ zipWith (\ dep (nm,ty) -> tyToTS ty dep) [0..] fnArgs

  let arg = zip (map fst fnArgs) argTypes
  let tco = isInfixOf "/*TCO*/" bodyStmt
  let bod = "{" ++ bodyStmt ++ "return " ++ bodyName ++ "; }"
  let fun = jsDefFun fnName arg tco bod
  let cur = jsDefCur fnName arg
  return $ fun ++ "\n" ++ cur

  where

  -- Generates top-level function
  jsDefFun name [] tco body = 
    let wrap = \x -> "(() => " ++ x ++ ")()"
        head = "const " ++ nameToJS name ++ "$ = "
    in head ++ wrap body
  jsDefFun name arg tco body =
    let loop = \ x -> concat ["{while(1)", x, "}"]
        head = "function " ++ nameToJS name ++ "$(" ++ intercalate "," (map (\ (nm,ty) -> nm++"/*:"++ty++"*/") arg) ++ ") "
    in head ++ (if tco then loop body else body)

  -- Generates top-level function (curried version)
  jsDefCur name arg =
    let head = "const " ++ nameToJS name ++ " = " ++ concat (map (\x -> x ++ " => ") (map fst arg))
        body = nameToJS name ++ "$" ++ (if null arg then "" else "(" ++ intercalate "," (map fst arg) ++ ")")
    in head ++ body

  -- Genreates a fresh name
  fresh :: ST.State Int String
  fresh = do
    n <- ST.get
    ST.put (n + 1)
    return $ "$x" ++ show n

  -- Assigns an expression to a name, or return it directly
  set :: String -> String -> ST.State Int String
  set name expr = return $ "var " ++ name ++ " = " ++ expr ++ ";"

  -- Compiles a name to JS
  nameToJS :: String -> String
  nameToJS x = "$" ++ map (\c -> if c == '/' || c == '.' || c == '-' || c == '#' then '$' else c) x

  -- Compiles an Oper to JS
  operToJS :: Oper -> String
  operToJS ADD = "+"
  operToJS SUB = "-"
  operToJS MUL = "*"
  operToJS DIV = "/"
  operToJS MOD = "%"
  operToJS EQ  = "==="
  operToJS NE  = "!=="
  operToJS LT  = "<"
  operToJS GT  = ">"
  operToJS LTE = "<="
  operToJS GTE = ">="
  operToJS AND = "&"
  operToJS OR  = "|"
  operToJS XOR = "^"
  operToJS LSH = "<<"
  operToJS RSH = ">>"

  -- Compiles a CType to TS
  tyToTS :: CT -> Int -> String
  tyToTS CSet dep =
    "Type"
  tyToTS CU64 dep =
    "BigInt"
  tyToTS CF64 dep =
    "Number"
  tyToTS (CADT cts) dep =
    intercalate " | " $ flip map cts $ \ (nm,fs) -> "{$:'" ++ nm ++ "'" ++ concat (map (\ (fn,ft) -> ", " ++ fn ++ ": " ++ tyToTS ft dep) fs) ++ "}"
  tyToTS (CMap typ) dep =
    "Map<BigInt, " ++ tyToTS typ dep ++ ">"
  tyToTS (CAll (nam,inp) bod) dep =
    let uid = nameToJS nam ++ "$" ++ show dep
    in "(" ++ uid ++ ":" ++ tyToTS inp dep ++ ") => " ++ tyToTS (bod (CVar uid dep)) (dep + 1)
  tyToTS (CRef nam) dep =
    nam
  tyToTS (CVar nam _) dep = 
    nam
  tyToTS (CApp fun arg) dep =
    tyToTS fun dep ++ "<" ++ tyToTS arg dep ++ ">"
  tyToTS CNul dep =
    "null"
  tyToTS term dep =
    "null"

  -- Compiles a CTerm to JS
  ctToJS :: Bool -> String -> CT -> Int -> ST.State Int String
  ctToJS tail var term dep = 
    -- trace ("COMPILE: " ++ showCT term 0) $
    go (red book term) where
    go CNul =
      set var "null"
    go CSet =
      set var "/*Type*/null"
    go ty@CU64 =
      set var $ "/*" ++ tyToTS ty dep ++ "*/null"
    go ty@CF64 =
      set var $ "/*" ++ tyToTS ty dep ++ "*/null"
    go ty@(CADT cts) = do
      set var $ "/*" ++ tyToTS ty dep ++ "*/null"
    go ty@(CMap typ) =
      set var $ "/*" ++ tyToTS ty dep ++ "*/null"
    go ty@(CAll (nam,inp) bod) =
      set var $ "/*" ++ tyToTS ty dep ++ "*/null"
    go tm@(CLam (nam,inp) bod) = do
      let (names, bodyTerm, _) = lams tm dep []
      bodyName <- fresh
      bodyStmt <- ctToJS False bodyName bodyTerm (dep + length names)
      set var $ concat ["(", intercalate " => " names, " => {", bodyStmt, "return ", bodyName, ";})"]
      where lams :: CT -> Int -> [String] -> ([String], CT, Maybe Term)
            lams (CLam (n,i) b) dep names =
              let uid = nameToJS n ++ "$" ++ show dep
              in lams (b (CVar uid dep)) (dep + 1) (uid : names)
            lams term dep names = (reverse names, term, Nothing)
    go app@(CApp fun arg) = do
      let (appFun, appArgs) = getAppChain app
      -- Tail Recursive Call
      if tail && isRecCall fnName (length fnArgs) appFun appArgs then do
        argDefs <- forM (zip (map fst fnArgs) appArgs) $ \ (paramName, appArgs) -> do
          argName <- fresh
          argStmt <- ctToJS False argName appArgs dep
          return (argStmt, paramName ++ " = " ++ argName ++ ";")
        let (argStmts, paramDefs) = unzip argDefs
        return $ concat argStmts ++ concat paramDefs ++ "/*TCO*/continue;"
      -- Saturated Call Optimization
      else if isSatCall book appFun appArgs then do
        let (CRef funName) = appFun
        argNamesStmts <- forM appArgs $ \arg -> do
          argName <- fresh
          argStmt <- ctToJS False argName arg dep
          return (argName, argStmt)
        retStmt <- set var $ concat [nameToJS funName, "$(", intercalate ", " (map fst argNamesStmts), ")"]
        return $ concat (map snd argNamesStmts ++ [retStmt])
      -- IO Actions
      else if isEffCall book appFun appArgs then do
        let (CHol name) = appFun
        case name of
          "IO_BIND" -> do
            let [_, _, call, cont] = appArgs
            callName <- fresh
            callStmt <- ctToJS False callName call dep
            contStmt <- ctToJS False var (CApp cont (CVar callName dep)) dep
            return $ concat [callStmt, contStmt]
          "IO_PURE" -> do
            let [_, value] = appArgs
            valueStmt <- ctToJS False var value dep
            return $ valueStmt
          "IO_SWAP" -> do
            let [key, val] = appArgs
            keyName  <- fresh
            keyStmt  <- ctToJS False keyName key dep
            valName  <- fresh
            valStmt  <- ctToJS False valName val dep
            resName  <- fresh
            resStmt  <- set resName (concat ["SWAP(", keyName, ", ", valName, ");"])
            doneStmt <- ctToJS False var (CVar resName 0) dep
            return $ concat [keyStmt, valStmt, resStmt, doneStmt]
          "IO_PRINT" -> do
            let [text] = appArgs
            textName <- fresh
            textStmt <- ctToJS False textName text dep
            doneStmt <- ctToJS False var (CCon "Unit" []) dep 
            return $ concat [textStmt, "console.log(LIST_TO_JSTR(", textName, "));", doneStmt]
          "IO_PROMPT" -> do
            error $ "TODO"
          "IO_READ" -> do
            let [path] = appArgs
            pathName <- fresh
            pathStmt <- ctToJS False pathName path dep
            let readStmt = concat
                  [ "try { var ", var, " = { $: 'Done', value: JSTR_TO_LIST(readFileSync(LIST_TO_JSTR(", pathName, "), 'utf8')) }; } "
                  , "catch (e) { var ", var, " = { $: 'Fail', error: e.message }; }"
                  ]
            return $ concat [pathStmt, readStmt]
          "IO_EXEC" -> do
            let [cmd] = appArgs
            cmdName  <- fresh
            cmdStmt  <- ctToJS False cmdName cmd dep
            retStmt  <- set var $ concat ["JSTR_TO_LIST(execSync(LIST_TO_JSTR(", cmdName, ")).toString())"]
            return $ concat [cmdStmt, retStmt]
          "IO_ARGS" -> do
            let [_] = appArgs
            retStmt  <- set var "JARRAY_TO_LIST(process.argv.slice(2), JSTR_TO_LIST)"
            return retStmt
          _ -> error $ "Unknown IO operation: " ++ name
      -- Normal Application
      else do
        funName <- fresh
        funStmt <- ctToJS False funName fun dep
        argName <- fresh
        argStmt <- ctToJS False argName arg dep
        retStmt <- set var $ concat ["(", funName, ")(", argName, ")"]
        return $ concat [funStmt, argStmt, retStmt]
    go (CCon nam fields) = do
      objStmt <- set var $ concat ["({$: \"", nam, "\"})"]
      setStmts <- forM fields $ \ (nm, tm) -> do
        fldName <- fresh
        fldStmt <- ctToJS False fldName tm dep
        setStmt <- return $ concat [var ++ "." ++ nm ++ " = " ++ fldName ++ ";"]
        return $ concat [fldStmt, setStmt]
      return $ concat $ [objStmt] ++ setStmts
    go (CMat val cses) = do
      let isRecord = length cses == 1 && not (any (\ (nm,_,_) -> nm == "_") cses)
      valName <- fresh
      valStmt <- ctToJS False valName val dep
      cases <- forM cses $ \ (cnam, fields, cbod) ->
        if cnam == "_" then do
          retStmt <- ctToJS tail var (CApp cbod (CVar valName 0)) dep
          return $ concat ["default: { " ++ retStmt, " break; }"]
        else do
          let bod = foldl CApp cbod (map (\ (fn,ft) -> (CVar (valName++"."++fn) 0)) fields)
          retStmt <- ctToJS tail var bod dep
          return $ if isRecord
            then retStmt
            else concat ["case \"", cnam, "\": { ", retStmt, " break; }"]
      let switch = if isRecord
            then concat [valStmt, unwords cases]
            else concat [valStmt, "switch (", valName, ".$) { ", unwords cases, " }"]
      return $ switch
    go (CSwi val zer suc) = do
      valName <- fresh
      valStmt <- ctToJS False valName val dep
      zerStmt <- ctToJS tail var zer dep
      sucStmt <- ctToJS tail var (CApp suc (COp2 CU64 SUB (CVar valName 0) (CNum 1))) dep
      let swiStmt = concat [valStmt, "if (", valName, " === 0n) { ", zerStmt, " } else { ", sucStmt, " }"]
      return $ swiStmt
    go (CKVs kvs def) = do
      dftStmt <- do
        dftName <- fresh
        dftStmt <- ctToJS False dftName def dep
        return $ concat [dftStmt, var, ".set(-1n, ", dftName, ");"]
      kvStmts <- forM (IM.toList kvs) $ \(k, v) -> do
        valName <- fresh
        valStmt <- ctToJS False valName v dep
        return $ concat [valStmt, var, ".set(", show k, "n, ", valName, ");"]
      let mapStmt = concat ["var ", var, " = new Map();", unwords kvStmts, dftStmt]
      return $ mapStmt
    go (CGet got nam map key bod) = do
      mapName <- fresh
      mapStmt <- ctToJS False mapName map dep
      keyName <- fresh
      keyStmt <- ctToJS False keyName key dep
      neoName <- fresh
      gotName <- fresh
      retStmt <- ctToJS tail var (bod (CVar gotName dep) (CVar neoName dep)) dep
      let gotStmt = concat ["var ", gotName, " = ", mapName, ".has(", keyName, ") ? ", mapName, ".get(", keyName, ") : ", mapName, ".get(-1n);"]
      let neoStmt = concat ["var ", neoName, " = ", mapName, ";"]
      return $ concat [mapStmt, keyStmt, gotStmt, neoStmt, retStmt]
    go (CPut got nam map key val bod) = do
      mapName <- fresh
      mapStmt <- ctToJS False mapName map dep
      keyName <- fresh
      keyStmt <- ctToJS False keyName key dep
      valName <- fresh
      valStmt <- ctToJS False valName val dep
      neoName <- fresh
      gotName <- fresh
      retStmt <- ctToJS tail var (bod (CVar gotName dep) (CVar neoName dep)) dep
      let gotStmt = concat ["var ", gotName, " = ", mapName, ".has(", keyName, ") ? ", mapName, ".get(", keyName, ") : ", mapName, ".get(-1n);"]
      let neoStmt = concat ["var ", neoName, " = ", mapName, "; ", mapName, ".set(", keyName, ", ", valName, ");"]
      return $ concat [mapStmt, keyStmt, valStmt, gotStmt, neoStmt, retStmt]
    go (CRef nam) =
      set var $ nameToJS nam
    go (CHol nam) =
      set var $ "null"
    go (CLet (nam,typ) val bod) = do
      let uid = nameToJS nam ++ "$" ++ show dep
      valStmt <- ctToJS False uid val dep
      bodStmt <- ctToJS tail var (bod (CVar uid dep)) (dep + 1)
      return $ concat [valStmt, bodStmt]
    go (CNum val) =
      set var $ show val ++ "n"
    go (CFlt val) =
      set var $ show val

    go (COp2 typ opr fst snd) = do
      let opr' = operToJS opr
      fstName <- fresh
      sndName <- fresh
      fstStmt <- ctToJS False fstName fst dep
      sndStmt <- ctToJS False sndName snd dep


      let retExpr = case typ of
            CF64 -> concat [fstName, " ", opr', " ", sndName]
            CU64 -> concat ["BigInt.asUintN(64, ", fstName, " ", opr', " ", sndName, ")"]
            _ -> error ("Invalid type for binary operation: " ++ showCT typ dep)
              
      retStmt <- set var retExpr
      return $ concat [fstStmt, sndStmt, retStmt]
    go (CLog msg nxt) = do
      msgName <- fresh
      msgStmt <- ctToJS False msgName msg dep
      nxtName <- fresh
      nxtStmt <- ctToJS tail nxtName nxt dep
      retStmt <- set var $ concat ["(console.log(LIST_TO_JSTR(", msgName, ")), ", nxtName, ")"]
      return $ concat [msgStmt, nxtStmt, retStmt]
    go (CVar nam _) =
      set var nam
    go (CTxt txt) =
      set var $ "JSTR_TO_LIST(`" ++ (concatMap (\c -> if c == '`' then "\\`" else [c]) txt) ++ "`)"
    go (CLst lst) =
      let cons = \x acc -> CCon "Cons" [("head", x), ("tail", acc)]
          nil  = CCon "Nil" []
      in  ctToJS False var (foldr cons nil lst) dep
    go (CNat val) =
      let succ = \x -> CCon "Succ" [("pred", x)]
          zero = CCon "Zero" []
      in  ctToJS False var (foldr (\_ acc -> succ acc) zero [1..val]) dep

prelude :: String
prelude = unlines [
  "import { readFileSync } from 'fs';",
  "import { execSync } from 'child_process';",
  "",
  "function LIST_TO_JSTR(list) {",
  "  try {",
  "    let result = '';",
  "    let current = list;",
  "    while (current.$ === 'Cons') {",
  "      result += String.fromCodePoint(Number(current.head));",
  "      current = current.tail;",
  "    }",
  "    if (current.$ === 'Nil') {",
  "      return result;",
  "    }",
  "  } catch (e) {}",
  "  return list;",
  "}",
  "",
  "function JSTR_TO_LIST(str) {",
  "  let list = {$: 'Nil'};",
  "  for (let i = str.length - 1; i >= 0; i--) {",
  "    list = {$: 'Cons', head: BigInt(str.charCodeAt(i)), tail: list};",
  "  }",
  "  return list;",
  "}",
  "",
  "function LIST_TO_JARRAY(list, decode) {",
  "  try {",
  "    let result = [];",
  "    let current = list;",
  "    while (current.$ === 'Cons') {",
  "      result += decode(current.head);",
  "      current = current.tail;",
  "    }",
  "    if (current.$ === 'Nil') {",
  "      return result;",
  "    }",
  "  } catch (e) {}",
  "  return list;",
  "}",
  "",
  "function JARRAY_TO_LIST(inp, encode) {",
  "  let out = {$: 'Nil'};",
  "  for (let i = inp.length - 1; i >= 0; i--) {",
  "    out = {$: 'Cons', head: encode(inp[i]), tail: out};",
  "  }",
  "  return out;",
  "}",
  "",
  "let MEMORY = new Map();",
  "function SWAP(key, val) {",
  "  var old = MEMORY.get(key) || 0n;",
  "  MEMORY.set(key, val);",
  "  return old;",
  "}"
  ]

generateJS :: CTBook -> (String, CT) -> String
generateJS book (name, ct) = ST.evalState (fnToJS book name ct) 0 ++ "\n\n"

defToCT :: Book -> (String, Term) -> (String, CT)
defToCT book (name, term) =
  case envRun (doAnnotate term) book of
    Done _ (term, fill) -> (name, termToCT book fill term Nothing 0)
    Fail _              -> error $ "COMPILATION_ERROR: " ++ name ++ " is ill-typed"

compileJS :: Book -> String
compileJS book =
  let ctDefs0 = flip map (topoSortBook book) (defToCT book)
      ctDefs1 = flip map ctDefs0 $ \ (nm,ct) -> (nm, removeUnreachables ct)
      ctDefs2 = flip map ctDefs1 $ \ (nm,ct) -> (nm, inline (M.fromList ctDefs1) ct)
      ctDefs3 = flip map ctDefs2 $ \ (nm,ct) -> (nm, liftLambdas ct 0)
      jsFns   = concatMap (generateJS (M.fromList ctDefs3)) ctDefs3
      exports = "export { " ++ intercalate ", " (getFunctionNames jsFns) ++ " }" 
      debug   = trace ("\nCompiled CTs:\n" ++ unlines (map (\(n,c) -> "- " ++ n ++ ":\n" ++ showCT c 0) ctDefs3))
  in prelude ++ "\n\n" ++ jsFns ++ "\n" ++ exports

-- Utils
-- -----

bindCT :: CT -> [(String,CT)] -> CT
bindCT CNul ctx = CNul
bindCT CSet ctx = CSet
bindCT CU64 ctx = CU64
bindCT CF64 ctx = CF64
bindCT (CADT cts) ctx =
  let cts' = map (\ (n,fs) -> (n, map (\ (fn,ft) -> (fn, bindCT ft ctx)) fs)) cts in
  CADT cts'
bindCT (CMap typ) ctx =
  CMap (bindCT typ ctx)
bindCT (CAll (nam,inp) bod) ctx =
  let inp' = bindCT inp ctx in
  let bod' = \x -> bindCT (bod (CVar nam 0)) ((nam, x) : ctx) in
  CAll (nam,inp') bod'
bindCT (CLam (nam,inp) bod) ctx =
  let inp' = bindCT inp ctx in
  let bod' = \x -> bindCT (bod (CVar nam 0)) ((nam, x) : ctx) in
  CLam (nam,inp') bod'
bindCT (CApp fun arg) ctx =
  let fun' = bindCT fun ctx in
  let arg' = bindCT arg ctx in
  CApp fun' arg'
bindCT (CCon nam arg) ctx =
  let arg' = map (\(f, x) -> (f, bindCT x ctx)) arg in
  CCon nam arg'
bindCT (CMat val cse) ctx =
  let val' = bindCT val ctx in
  let cse' = map (\(cn,fs,cb) -> (cn, fs, bindCT cb ctx)) cse in
  CMat val' cse'
bindCT (CRef nam) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> CRef nam
bindCT (CHol nam) ctx =
  CHol nam
bindCT (CLet (nam,typ) val bod) ctx =
  let typ' = bindCT typ ctx in
  let val' = bindCT val ctx in
  let bod' = \x -> bindCT (bod (CVar nam 0)) ((nam, x) : ctx) in
  CLet (nam,typ') val' bod'
bindCT (CNum val) ctx = CNum val
bindCT (CFlt val) ctx = CFlt val
bindCT (COp2 typ opr fst snd) ctx =
  let fst' = bindCT fst ctx in
  let snd' = bindCT snd ctx in
  let typ' = bindCT typ ctx in
  COp2 typ' opr fst' snd'
bindCT (CSwi val zer suc) ctx =
  let val' = bindCT val ctx in
  let zer' = bindCT zer ctx in
  let suc' = bindCT suc ctx in
  CSwi val' zer' suc'
bindCT (CKVs kvs def) ctx =
  let kvs' = IM.map (\v -> bindCT v ctx) kvs in
  let def' = bindCT def ctx in
  CKVs kvs' def'
bindCT (CGet got nam map key bod) ctx =
  let map' = bindCT map ctx in
  let key' = bindCT key ctx in
  let bod' = \x y -> bindCT (bod (CVar got 0) (CVar nam 0)) ((nam, y) : (got, x) : ctx) in
  CGet got nam map' key' bod'
bindCT (CPut got nam map key val bod) ctx =
  let map' = bindCT map ctx in
  let key' = bindCT key ctx in
  let val' = bindCT val ctx in
  let bod' = \x y -> bindCT (bod (CVar got 0) (CVar nam 0)) ((nam, y) : (got, x) : ctx) in
  CPut got nam map' key' val' bod'
bindCT (CLog msg nxt) ctx =
  let msg' = bindCT msg ctx in
  let nxt' = bindCT nxt ctx in
  CLog msg' nxt'
bindCT (CVar nam idx) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> CVar nam idx
bindCT (CTxt txt) ctx = CTxt txt
bindCT (CLst lst) ctx =
  let lst' = map (\x -> bindCT x ctx) lst in
  CLst lst'
bindCT (CNat val) ctx = CNat val

rnCT :: CT -> [(String,CT)] -> CT
rnCT CNul ctx = CNul
rnCT CSet ctx = CSet
rnCT CU64 ctx = CU64
rnCT CF64 ctx = CF64
rnCT (CADT cts) ctx =
  let cts' = map (\ (n,fs) -> (n, map (\ (fn,ft) -> (fn, rnCT ft ctx)) fs)) cts in
  CADT cts'
rnCT (CMap typ) ctx =
  let typ' = rnCT typ ctx
  in (CMap typ')
rnCT (CAll (nam,inp) bod) ctx =
  let nam' = "x" ++ show (length ctx) in
  let inp' = rnCT inp ctx in
  let bod' = \x -> rnCT (bod (CVar nam' 0)) ((nam', x) : ctx) in
  CAll (nam',inp') bod'
rnCT (CLam (nam,inp) bod) ctx =
  let nam' = "x" ++ show (length ctx) in
  let inp' = rnCT inp ctx in
  let bod' = \x -> rnCT (bod (CVar nam' 0)) ((nam', x) : ctx) in
  CLam (nam',inp') bod'
rnCT (CApp fun arg) ctx =
  let fun' = rnCT fun ctx in
  let arg' = rnCT arg ctx in
  CApp fun' arg'
rnCT (CCon nam arg) ctx =
  let arg' = map (\(f, x) -> (f, rnCT x ctx)) arg in
  CCon nam arg'
rnCT (CMat val cse) ctx =
  let val' = rnCT val ctx in
  let cse' = map (\(cn,fs,cb) -> (cn, fs, rnCT cb ctx)) cse in
  CMat val' cse'
rnCT (CRef nam) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> CRef nam
rnCT (CLet (nam,typ) val bod) ctx =
  let typ' = rnCT typ ctx in
  let val' = rnCT val ctx in
  let bod' = \x -> rnCT (bod (CVar nam 0)) ((nam, x) : ctx) in
  CLet (nam,typ') val' bod'
rnCT (CNum val) ctx = CNum val
rnCT (CFlt val) ctx = CFlt val
rnCT (COp2 typ opr fst snd) ctx =
  let fst' = rnCT fst ctx in
  let snd' = rnCT snd ctx in
  let typ' = rnCT typ ctx in
  COp2 typ' opr fst' snd'
rnCT (CSwi val zer suc) ctx =
  let val' = rnCT val ctx in
  let zer' = rnCT zer ctx in
  let suc' = rnCT suc ctx in
  CSwi val' zer' suc'
rnCT (CKVs kvs def) ctx =
  let kvs' = IM.map (\v -> rnCT v ctx) kvs in
  let def' = rnCT def ctx in
  CKVs kvs' def'
rnCT (CGet got nam map key bod) ctx =
  let map' = rnCT map ctx in
  let key' = rnCT key ctx in
  let bod' = \x y -> rnCT (bod (CVar got 0) (CVar nam 0)) ((got, x) : (nam, y) : ctx) in
  CGet got nam map' key' bod'
rnCT (CPut got nam map key val bod) ctx =
  let map' = rnCT map ctx in
  let key' = rnCT key ctx in
  let val' = rnCT val ctx in
  let bod' = \x y -> rnCT (bod (CVar got 0) (CVar nam 0)) ((got, x) : (nam, y) : ctx) in
  CPut got nam map' key' val' bod'
rnCT (CLog msg nxt) ctx =
  let msg' = rnCT msg ctx in
  let nxt' = rnCT nxt ctx in
  CLog msg' nxt'
rnCT (CVar nam idx) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> CVar nam idx
rnCT (CTxt txt) ctx = CTxt txt
rnCT (CLst lst) ctx =
  let lst' = map (\x -> rnCT x ctx) lst in
  CLst lst'
rnCT (CNat val) ctx = CNat val

getAppChain :: CT -> (CT, [CT])
getAppChain (CApp fun arg) =
  let (f, args) = getAppChain fun
  in (f, args ++ [arg])
getAppChain term = (term, [])

isNul :: CT -> Bool
isNul CNul = True
isNul _    = False

getFunctionNames :: String -> [String]
getFunctionNames js = 
  [ name | line <- lines js,
           "const " `isPrefixOf` line,
           let parts = words line,
           length parts >= 2,
           let name = head $ words $ parts !! 1,
           not $ "$" `isSuffixOf` name  -- Skip internal functions ending with $
  ]

-- Stringification
-- ---------------

-- TODO: implement a showCT :: CT -> String function
showCT :: CT -> Int -> String
showCT CNul                     dep = "*"
showCT CSet                     dep = "Set"
showCT CU64                     dep = "U64"
showCT CF64                     dep = "F64"
showCT (CADT cts)               dep = "data{" ++ concatMap (\ (n,fs) -> "#" ++ n ++ " " ++ concatMap (\ (fn,ft) -> fn ++ ":" ++ showCT ft dep ++ " ") fs) cts ++ "}"
showCT (CMap typ)               dep = "(Map " ++ showCT typ dep ++ ")"
showCT (CLam (nam,inp) bod)     dep = "λ(" ++ nam ++ ": " ++ showCT inp dep ++ "). " ++ showCT (bod (CVar nam dep)) (dep+1)
showCT (CAll (nam,inp) bod)     dep = "∀(" ++ nam ++ ": " ++ showCT inp dep ++ "). " ++ showCT (bod (CVar nam dep)) (dep+1)
showCT (CApp fun arg)           dep = "(" ++ showCT fun dep ++ " " ++ showCT arg dep ++ ")"
showCT (CCon nam fields)        dep = "#" ++ nam ++ "{" ++ concatMap (\ (f,v) -> f ++ ":" ++ showCT v dep ++ " ") fields ++ "}"
showCT (CMat val cses)          dep = "match " ++ showCT val dep ++ " {" ++ concatMap (\(cn,fs,cb) -> "#" ++ cn ++ ":" ++ showCT cb dep ++ " ") cses ++ "}"
showCT (CRef nam)               dep = nam
showCT (CHol nam)               dep = nam
showCT (CLet (nam,typ) val bod) dep = "let " ++ nam ++ " : " ++ showCT typ dep ++ " = " ++ showCT val dep ++ "; " ++ showCT (bod (CVar nam dep)) (dep+1)
showCT (CNum val)               dep = show val
showCT (CFlt val)               dep = show val
showCT (COp2 typ opr fst snd)   dep = "(<op> " ++ showCT fst dep ++ " " ++ showCT snd dep ++ ")"
showCT (CSwi val zer suc)       dep = "switch " ++ showCT val dep ++ " {0:" ++ showCT zer dep ++ " _: " ++ showCT suc dep ++ "}"
showCT (CKVs kvs def)           dep = "{" ++ unwords (map (\(k,v) -> show k ++ ":" ++ showCT v dep) (IM.toList kvs)) ++ " | " ++ showCT def dep ++ "}"
showCT (CGet g n m k b)         dep = "get " ++ g ++ " = " ++ n ++ "@" ++ showCT m dep ++ "[" ++ showCT k dep ++ "] " ++ showCT (b (CVar g dep) (CVar n dep)) (dep+2)
showCT (CPut g n m k v b)       dep = "put " ++ g ++ " = " ++ n ++ "@" ++ showCT m dep ++ "[" ++ showCT k dep ++ "] := " ++ showCT v dep ++ " " ++ showCT (b (CVar g dep) (CVar n dep)) (dep+2)
showCT (CLog msg nxt)           dep = "log(" ++ showCT msg dep ++ "," ++ showCT nxt dep ++ ")"
showCT (CVar nam dep)           _   = nam ++ "^" ++ show dep
showCT (CTxt txt)               dep = show txt
showCT (CLst lst)               dep = "[" ++ unwords (map (\x -> showCT x dep) lst) ++ "]"
showCT (CNat val)               dep = show val


================================================
FILE: src/Kind/Env.hs
================================================
module Kind.Env where

import Kind.Type

import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M

-- Environment
-- -----------

envBind :: Env a -> (a -> Env b) -> Env b
envBind (Env a) b = Env $ \state -> case a state of
  Done state' value -> let Env b' = b value in b' state'
  Fail state'       -> Fail state'

envPure :: a -> Env a
envPure a = Env $ \state -> Done state a

envFail :: Env a
envFail = Env $ \state -> Fail state

envRun :: Env a -> Book -> Res a
envRun (Env chk) book = chk (State book IM.empty [] [])

envLog :: Info -> Env Int
envLog log = Env $ \ (State book fill susp logs) -> Done (State book fill susp (log : logs)) 1

envSnapshot :: Env State
envSnapshot = Env $ \state -> Done state state

envRewind :: State -> Env Int
envRewind state = Env $ \_ -> Done state 0

envSusp :: Check -> Env ()
envSusp chk = Env $ \ (State book fill susp logs) -> Done (State book fill (susp ++ [chk]) logs) ()

envFill :: Int -> Term -> Env ()
envFill k v = Env $ \ (State book fill susp logs) -> Done (State book (IM.insert k v fill) susp logs) ()

envGetFill :: Env Fill
envGetFill = Env $ \ (State book fill susp logs) -> Done (State book fill susp logs) fill

envGetBook :: Env Book
envGetBook = Env $ \ (State book fill susp logs) -> Done (State book fill susp logs) book

envTakeSusp :: Env [Check]
envTakeSusp = Env $ \ (State book fill susp logs) -> Done (State book fill [] logs) susp

instance Functor Env where
  fmap f (Env chk) = Env $ \logs -> case chk logs of
    Done logs' a -> Done logs' (f a)
    Fail logs' -> Fail logs'

instance Applicative Env where
  pure = envPure
  (Env chkF) <*> (Env chkA) = Env $ \logs -> case chkF logs of
    Done logs' f -> case chkA logs' of
      Done logs'' a -> Done logs'' (f a)
      Fail logs'' -> Fail logs''
    Fail logs' -> Fail logs'

instance Monad Env where
  (Env a) >>= b = envBind (Env a) b


================================================
FILE: src/Kind/Equal.hs
================================================
-- //./Type.hs//

module Kind.Equal where

import Control.Monad (zipWithM)

import Debug.Trace

import Kind.Type
import Kind.Env
import Kind.Reduce
import Kind.Show

import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as IM

-- Equality
-- --------

-- Checks if two terms are equal, after reduction steps
equal :: Term -> Term -> Int -> Env Bool
equal a b dep = debug ("== " ++ showTermGo False a dep ++ "\n.. " ++ showTermGo False b dep) $ do
  -- If both terms are identical, return true
  state <- envSnapshot
  is_id <- identical a b dep
  if is_id then do
    envPure True
  -- Otherwise, reduces both terms to wnf
  else do
    envRewind state
    book <- envGetBook
    fill <- envGetFill
    let aWnf = reduce book fill 2 a
    let bWnf = reduce book fill 2 b
    -- If both term wnfs are identical, return true
    state <- envSnapshot
    is_id <- identical aWnf bWnf dep
    if is_id then do
      envPure True
    -- Otherwise, check if they're component-wise equal
    else do
      envRewind state
      similar aWnf bWnf dep

-- Checks if two terms are already syntactically identical
identical :: Term -> Term -> Int -> Env Bool
identical a b dep = do
  fill <- envGetFill
  debug ("ID " ++ showTermGo False a dep ++ "\n.. " ++ showTermGo False b dep ++ "\n" ++ (unlines $ map (\(k,v) -> "~" ++ show k ++ " = " ++ showTermGo False v dep) $ IM.toList fill)) $ go a b dep
 where
  go (All aNam aInp aBod) (All bNam bInp bBod) dep = do
    iInp <- identical aInp bInp dep
    iBod <- identical (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
    return (iInp && iBod)
  go (Lam aNam aBod) (Lam bNam bBod) dep =
    identical (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
  go (App aFun aArg) (App bFun bArg) dep = do
    iFun <- identical aFun bFun dep
    iArg <- identical aArg bArg dep
    return (iFun && iArg)
  go (Slf aNam aTyp aBod) (Slf bNam bTyp bBod) dep =
    identical aTyp bTyp dep
  go (Ins aVal) b dep =
    identical aVal b dep
  go a (Ins bVal) dep =
    identical a bVal dep
  go (ADT aScp aCts aTyp) (ADT bScp bCts bTyp) dep = do
    identical aTyp bTyp dep
  go (Con aNam aArg) (Con bNam bArg) dep = do
    if aNam == bNam && length aArg == length bArg
      then and <$> zipWithM (\(_, aVal) (_, bVal) -> identical aVal bVal dep) aArg bArg
      else return False
  go (Mat aCse) (Mat bCse) dep = do
    if length aCse == length bCse
      then and <$> zipWithM goCse aCse bCse
      else return False
  go (Let aNam aVal aBod) b dep =
    identical (aBod aVal) b dep
  go a (Let bNam bVal bBod) dep =
    identical a (bBod bVal) dep
  go (Use aNam aVal aBod) b dep =
    identical (aBod aVal) b dep
  go a (Use bNam bVal bBod) dep =
    identical a (bBod bVal) dep
  go Set Set dep =
    return True
  go (Ann chk aVal aTyp) b dep =
    identical aVal b dep
  go a (Ann chk bVal bTyp) dep =
    identical a bVal dep
  go (Met aUid aSpn) b dep = do
    fill <- envGetFill
    case IM.lookup aUid fill of
      Just sol -> identical sol b dep
      Nothing  -> unify aUid aSpn b dep
  go a (Met bUid bSpn) dep = do
    fill <- envGetFill
    case IM.lookup bUid fill of
      Just sol -> identical a sol dep
      Nothing  -> unify bUid bSpn a dep
  go (Log aMsg aNxt) b dep =
    identical aNxt b dep
  go a (Log bMsg bNxt) dep =
    identical a bNxt dep
  go (Hol aNam aCtx) b dep =
    return True
  go a (Hol bNam bCtx) dep =
    return True
  go U64 U64 dep =
    return True
  go F64 F64 dep =
    return True
  go (Num aVal) (Num bVal) dep =
    return (aVal == bVal)
  go (Flt aVal) (Flt bVal) dep =
    return (aVal == bVal)
  go (Op2 aOpr aFst aSnd) (Op2 bOpr bFst bSnd) dep = do
    iFst <- identical aFst bFst dep
    iSnd <- identical aSnd bSnd dep
    return (iFst && iSnd)
  go (Swi aZer aSuc) (Swi bZer bSuc) dep = do
    iZer <- identical aZer bZer dep
    iSuc <- identical aSuc bSuc dep
    return (iZer && iSuc)
  go (Map aTyp) (Map bTyp) dep =
    identical aTyp bTyp dep
  go (KVs aMap aDef) (KVs bMap bDef) dep = do
    iDef <- identical aDef bDef dep
    iMap <- flip mapM (IM.toList aMap) $ \ (aKey,aVal) ->
      case IM.lookup aKey bMap of
        Just bVal -> identical aVal bVal dep
        Nothing   -> return False
    return (iDef && and iMap && IM.size aMap == IM.size bMap)
  go (Get aGot aNam aMap aKey aBod) (Get bGot bNam bMap bKey bBod) dep = do
    iMap <- identical aMap bMap dep
    iKey <- identical aKey bKey dep
    iBod <- identical (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
    return (iMap && iKey && iBod)
  go (Put aGot aNam aMap aKey aVal aBod) (Put bGot bNam bMap bKey bVal bBod) dep = do
    iMap <- identical aMap bMap dep
    iKey <- identical aKey bKey dep
    iVal <- identical aVal bVal dep
    iBod <- identical (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
    return (iMap && iKey && iVal && iBod)
  go (Txt aTxt) (Txt bTxt) dep =
    return (aTxt == bTxt)
  go (Lst aLst) (Lst bLst) dep =
    if length aLst == length bLst
      then and <$> zipWithM (\a b -> identical a b dep) aLst bLst
      else return False
  go (Nat aVal) (Nat bVal) dep =
    return (aVal == bVal)
  go (Src aSrc aVal) b dep =
    identical aVal b dep
  go a (Src bSrc bVal) dep =
    identical a bVal dep
  go (Ref aNam) (Ref bNam) dep =
    return (aNam == bNam)
  go (Var aNam aIdx) (Var bNam bIdx) dep =
    return (aIdx == bIdx)
  go a b dep =
    return False

  goCtr (Ctr aCNm aTele) (Ctr bCNm bTele) = do
    if aCNm == bCNm
      then goTele aTele bTele dep
      else return False

  goCse (aCNam, aCBod) (bCNam, bCBod) = do
    if aCNam == bCNam
      then identical aCBod bCBod dep
      else return False

  goTele :: Tele -> Tele -> Int -> Env Bool
  goTele (TRet aTerm) (TRet bTerm) dep = identical aTerm bTerm dep
  goTele (TExt aNam aTyp aBod) (TExt bNam bTyp bBod) dep = do
    iTyp <- identical aTyp bTyp dep
    iBod <- goTele (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
    return (iTyp && iBod)
  goTele _ _ _ = return False

-- Checks if two terms are component-wise equal
similar :: Term -> Term -> Int -> Env Bool
similar a b dep = go a b dep where
  go (All aNam aInp aBod) (All bNam bInp bBod) dep = do
    eInp <- equal aInp bInp dep
    eBod <- equal (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
    return (eInp && eBod)
  go (Lam aNam aBod) (Lam bNam bBod) dep =
    equal (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
  go (App aFun aArg) (App bFun bArg) dep = do
    eFun <- similar aFun bFun dep
    eArg <- equal aArg bArg dep
    return (eFun && eArg)
  go (Slf aNam aTyp aBod) (Slf bNam bTyp bBod) dep = do
    book <- envGetBook
    similar (reduce book IM.empty 0 aTyp) (reduce book IM.empty 0 bTyp) dep
  go (ADT aScp aCts aTyp) (ADT bScp bCts bTyp) dep = do
    book <- envGetBook
    similar (reduce book IM.empty 0 aTyp) (reduce book IM.empty 0 bTyp) dep
    -- eSlf <- zipWithM (\ax bx -> equal ax bx dep) aScp bScp
    -- if and eSlf && length aCts == length bCts
      -- then and <$> zipWithM goCtr aCts bCts
      -- else return False
  go (Con aNam aArg) (Con bNam bArg) dep = do
    if aNam == bNam && length aArg == length bArg
      then and <$> zipWithM (\(_, aVal) (_, bVal) -> equal aVal bVal dep) aArg bArg
      else return False
  go (Mat aCse) (Mat bCse) dep = do
    if length aCse == length bCse
      then and <$> zipWithM goCse aCse bCse
      else return False
  go (Op2 aOpr aFst aSnd) (Op2 bOpr bFst bSnd) dep = do
    eFst <- equal aFst bFst dep
    eSnd <- equal aSnd bSnd dep
    return (eFst && eSnd)
  go (Swi aZer aSuc) (Swi bZer bSuc) dep = do
    eZer <- equal aZer bZer dep
    eSuc <- equal aSuc bSuc dep
    return (eZer && eSuc)
  go (Map aTyp) (Map bTyp) dep = do
    equal aTyp bTyp dep
  go (KVs aMap aDef) (KVs bMap bDef) dep = do
    eDef <- equal aDef bDef dep
    eMap <- flip mapM (IM.toList aMap) $ \ (aKey,aVal) ->
      case IM.lookup aKey bMap of
        Just bVal -> equal aVal bVal dep
        Nothing   -> return False
    return (eDef && and eMap && IM.size aMap == IM.size bMap)
  go (Get aGot aNam aMap aKey aBod) (Get bGot bNam bMap bKey bBod) dep = do
    eMap <- equal aMap bMap dep
    eKey <- equal aKey bKey dep
    eBod <- equal (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
    return (eMap && eKey && eBod)
  go (Put aGot aNam aMap aKey aVal aBod) (Put bGot bNam bMap bKey bVal bBod) dep = do
    eMap <- equal aMap bMap dep
    eKey <- equal aKey bKey dep
    eVal <- equal aVal bVal dep
    eBod <- equal (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
    return (eMap && eKey && eVal && eBod)
  go a b dep = identical a b dep

  goCtr (Ctr aCNm aTele) (Ctr bCNm bTele) = do
    if aCNm == bCNm
      then goTele aTele bTele dep
      else return False

  goCse (aCNam, aCBod) (bCNam, bCBod) = do
    if aCNam == bCNam
      then equal aCBod bCBod dep
      else return False

  goTele :: Tele -> Tele -> Int -> Env Bool
  goTele (TRet aTerm) (TRet bTerm) dep = equal aTerm bTerm dep
  goTele (TExt aNam aTyp aBod) (TExt bNam bTyp bBod) dep = do
    eTyp <- equal aTyp bTyp dep
    eBod <- goTele (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
    return (eTyp && eBod)
  goTele _ _ _ = return False

-- Unification
-- -----------

-- If possible, solves a (?X x y z ...) = K problem, generating a subst.
unify :: Int -> [Term] -> Term -> Int -> Env Bool
unify uid spn b dep = do
  book <- envGetBook
  fill <- envGetFill

  -- is this hole not already solved?
  let solved = IM.member uid fill

  -- does the spine satisfies conditions?
  let solvable = valid fill spn []

  -- is the solution not recursive?
  let no_loops = not $ occur book fill uid b dep

  debug ("unify: " ++ show uid ++ " " ++ showTermGo False b dep ++ " | " ++ show solved ++ " " ++ show solvable ++ " " ++ show no_loops) $ do
    if not solved && solvable && no_loops then do
      let solution = solve book fill uid spn b
      debug ("solve: " ++ show uid ++ " " ++ showTermGo False solution dep ++ " | spn: " ++ show (map (\t -> showTermGo False t dep) spn)) $ envFill uid solution
      return True

    -- Otherwise, return true iff both are identical metavars
    else case b of
      (Src bSrc bVal) -> unify uid spn bVal dep
      (Met bUid bSpn) -> return $ uid == bUid
      other           -> return $ False

-- Checks if a problem is solveable by pattern unification.
valid :: Fill -> [Term] -> [Int] -> Bool
valid fill []        vars = True
valid fill (x : spn) vars = case reduce M.empty fill 0 x of
  (Var nam idx) -> not (elem idx vars) && valid fill spn (idx : vars)
  otherwise     -> False

-- Generates the solution, adding binders and renaming variables.
solve :: Book -> Fill -> Int -> [Term] -> Term -> Term
solve book fill uid []        b = b
solve book fill uid (x : spn) b = case reduce book fill 0 x of
  (Var nam idx) -> Lam nam $ \x -> subst idx x (solve book fill uid spn b)
  otherwise     -> error "unreachable"

-- Checks if a metavar uid occurs recursively inside a term
occur :: Book -> Fill -> Int -> Term -> Int -> Bool
occur book fill uid term dep = go term dep where
  go (All nam inp bod) dep =
    let o_inp = go inp dep
        o_bod = go (bod (Var nam dep)) (dep + 1)
    in o_inp || o_bod
  go (Lam nam bod) dep =
    let o_bod = go (bod (Var nam dep)) (dep + 1)
    in  o_bod
  go (App fun arg) dep =
    let o_fun = go fun dep
        o_arg = go arg dep
    in o_fun || o_arg
  go (Ann chk val typ) dep =
    let o_val = go val dep
        o_typ = go typ dep
    in o_val || o_typ
  go (Slf nam typ bod) dep =
    let o_typ = go typ dep
        o_bod = go (bod (Var nam dep)) (dep + 1)
    in o_typ || o_bod
  go (Ins val) dep =
    let o_val = go val dep
    in o_val
  go (ADT scp cts typ) dep =
    let o_scp = any (\x -> go x dep) scp
        o_cts = any (\(Ctr _ tele) -> goTele tele dep) cts
        a_typ = go typ dep
    in o_scp || o_cts || a_typ
  go (Con nam arg) dep =
    any (\(_, x) -> go x dep) arg
  go (Mat cse) dep =
    any (\ (_, cbod) -> go cbod dep) cse
  go (Let nam val bod) dep =
    let o_val = go val dep
        o_bod = go (bod (Var nam dep)) (dep + 1)
    in o_val || o_bod
  go (Use nam val bod) dep =
    let o_val = go val dep
        o_bod = go (bod (Var nam dep)) (dep + 1)
    in o_val || o_bod
  go (Log msg nxt) dep =
    let o_msg = go msg dep
        o_nxt = go nxt dep
    in o_msg || o_nxt
  go (Hol nam ctx) dep =
    False
  go (Op2 opr fst snd) dep =
    let o_fst = go fst dep
        o_snd = go snd dep
    in o_fst || o_snd
  go (Swi zer suc) dep =
    let o_zer = go zer dep
        o_suc = go suc dep
    in o_zer || o_suc
  go (Map typ) dep =
    let o_typ = go typ dep
    in o_typ
  go (KVs map def) dep =
    let o_map = any (\(_, x) -> go x dep) (IM.toList map)
        o_def = go def dep
    in o_map || o_def
  go (Get got nam map key bod) dep =
    let o_map = go map dep
        o_key = go key dep
        o_bod = go (bod (Var got dep) (Var nam dep)) (dep + 2)
    in o_map || o_key || o_bod
  go (Put got nam map key val bod) dep =
    let o_map = go map dep
        o_key = go key dep
        o_val = go val dep
        o_bod = go (bod (Var got dep) (Var nam dep)) (dep + 2)
    in o_map || o_key || o_val || o_bod
  go (Src src val) dep =
    let o_val = go val dep
    in o_val
  go (Met bUid bSpn) dep =
    case reduce book fill 2 (Met bUid bSpn) of
      Met bUid bSpn -> uid == bUid
      term          -> go term dep
  go _ dep =
    False

  goTele :: Tele -> Int -> Bool
  goTele (TRet term) dep = go term dep
  goTele (TExt nam typ bod) dep =
    let o_typ = go typ dep
        o_bod = goTele (bod (Var nam dep)) (dep + 1)
    in o_typ || o_bod

-- Substitution
-- ------------

-- This is the ugly / slow part of Kind. See: https://gist.github.com/VictorTaelin/48eed41a8eca3500721c06dfec72d48c

-- Behaves like 'identical', except it is pure and returns a Bool.
same :: Term -> Term -> Int -> Bool
same (All aNam aInp aBod) (All bNam bInp bBod) dep =
  let sInp = same aInp bInp dep
      sBod = same (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
  in sInp && sBod
same (Lam aNam aBod) (Lam bNam bBod) dep =
  let sBod = same (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
  in  sBod
same (App aFun aArg) (App bFun bArg) dep =
  let sFun = same aFun bFun dep
      sArg = same aArg bArg dep
  in sFun && sArg
same (Slf aNam aTyp aBod) (Slf bNam bTyp bBod) dep =
  let sTyp = same aTyp bTyp dep
  in  sTyp
same (Ins aVal) b dep =
  same aVal b dep
same a (Ins bVal) dep =
  same a bVal dep
same (ADT aScp aCts aTyp) (ADT bScp bCts bTyp) dep =
  -- let sSlf = and $ zipWith (\ax bx -> same ax bx dep) aScp bScp
      -- sCts = length aCts == length bCts && and (zipWith (\ a b -> sameCtr a b dep) aCts bCts)
  let sTyp = same aTyp bTyp dep
  in sTyp
same (Con aNam aArg) (Con bNam bArg) dep =
  let sNam = aNam == bNam
      sArg = length aArg == length bArg && and (zipWith (\(_, aVal) (_, bVal) -> same aVal bVal dep) aArg bArg)
  in sNam && sArg
same (Mat aCse) (Mat bCse) dep =
  let sCse = length aCse == length bCse && and (zipWith (\ a b -> sameCse a b dep) aCse bCse)
  in  sCse
same (Let aNam aVal aBod) b dep =
  same (aBod aVal) b dep
same a (Let bNam bVal bBod) dep =
  same a (bBod bVal) dep
same (Use aNam aVal aBod) b dep =
  same (aBod aVal) b dep
same a (Use bNam bVal bBod) dep =
  same a (bBod bVal) dep
same Set Set dep =
  True
same (Ann chk aVal aTyp) b dep =
  same aVal b dep
same a (Ann chk bVal bTyp) dep =
  same a bVal dep
same (Met aUid aSpn) b dep =
  False
same a (Met bUid bSpn) dep =
  False
-- TODO: Log
same (Log aMsg aNxt) b dep =
  same aNxt b dep
same a (Log bMsg bNxt) dep =
  same a bNxt dep
same (Hol aNam aCtx) b dep =
  True
same a (Hol bNam bCtx) dep =
  True
same U64 U64 dep =
  True
same F64 F64 dep =
  True
same (Num aVal) (Num bVal) dep =
  aVal == bVal
same (Flt aVal) (Flt bVal) dep =
  aVal == bVal
same (Op2 aOpr aFst aSnd) (Op2 bOpr bFst bSnd) dep =
  same aFst bFst dep && same aSnd bSnd dep
same (Swi aZer aSuc) (Swi bZer bSuc) dep =
  same aZer bZer dep && same aSuc bSuc dep
same (Map aTyp) (Map bTyp) dep =
  same aTyp bTyp dep
same (KVs aMap aDef) (KVs bMap bDef) dep =
  let sDef = same aDef bDef dep
      sMap = IM.size aMap == IM.size bMap && and (map (\ (aKey,aVal) -> maybe False (\bVal -> same aVal bVal dep) (IM.lookup aKey bMap)) (IM.toList aMap))
  in sDef && sMap
same (Get aGot aNam aMap aKey aBod) (Get bGot bNam bMap bKey bBod) dep =
  let sMap = same aMap bMap dep
      sKey = same aKey bKey dep
      sBod = same (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
  in sMap && sKey && sBod
same (Put aGot aNam aMap aKey aVal aBod) (Put bGot bNam bMap bKey bVal bBod) dep =
  let sMap = same aMap bMap dep
      sKey = same aKey bKey dep
      sVal = same aVal bVal dep
      sBod = same (aBod (Var aGot dep) (Var aNam dep)) (bBod (Var bGot dep) (Var bNam dep)) (dep + 2)
  in sMap && sKey && sVal && sBod
same (Txt aTxt) (Txt bTxt) dep =
  aTxt == bTxt
same (Lst aLst) (Lst bLst) dep =
  length aLst == length bLst && and (zipWith (\a b -> same a b dep) aLst bLst)
same (Nat aVal) (Nat bVal) dep =
  aVal == bVal
same (Src aSrc aVal) b dep =
  same aVal b dep
same a (Src bSrc bVal) dep =
  same a bVal dep
same (Ref aNam) (Ref bNam) dep =
  aNam == bNam
same (Var aNam aIdx) (Var bNam bIdx) dep =
  aIdx == bIdx
same _ _ _ = False

-- Auxiliary functions
sameCtr :: Ctr -> Ctr -> Int -> Bool
sameCtr (Ctr aCNm aTele) (Ctr bCNm bTele) dep =
  if aCNm == bCNm
    then sameTele aTele bTele dep
    else False

sameCse :: (String, Term) -> (String, Term) -> Int -> Bool
sameCse (aCNam, aCBod) (bCNam, bCBod) dep =
  if aCNam == bCNam
    then same aCBod bCBod dep
    else False

sameTele :: Tele -> Tele -> Int -> Bool
sameTele (TRet aTerm) (TRet bTerm) dep = same aTerm bTerm dep
sameTele (TExt aNam aTyp aBod) (TExt bNam bTyp bBod) dep =
  let sTyp = same aTyp bTyp dep
      sBod = sameTele (aBod (Var aNam dep)) (bBod (Var bNam dep)) (dep + 1)
  in sTyp && sBod
sameTele _ _ _ = False

-- Substitutes a Bruijn level variable by a neo value in term.
subst :: Int -> Term -> Term -> Term
subst lvl neo term = go term where
  go (All nam inp bod) = All nam (go inp) (\x -> go (bod (Sub x)))
  go (Lam nam bod)     = Lam nam (\x -> go (bod (Sub x)))
  go (App fun arg)     = App (go fun) (go arg)
  go (Ann chk val typ) = Ann chk (go val) (go typ)
  go (Slf nam typ bod) = Slf nam (go typ) (\x -> go (bod (Sub x)))
  go (Ins val)         = Ins (go val)
  go (ADT scp cts typ) = ADT (map go scp) (map goCtr cts) (go typ)
  go (Con nam arg)     = Con nam (map (\(f, t) -> (f, go t)) arg)
  go (Mat cse)         = Mat (map goCse cse)
  go (Swi zer suc)     = Swi (go zer) (go suc)
  go (Map typ)         = Map (go typ)
  go (KVs map def)     = KVs (IM.map go map) (go def)
  go (Get g n m k b)   = Get g n (go m) (go k) (\x y -> go (b x y))
  go (Put g n m k v b) = Put g n (go m) (go k) (go v) (\x y -> go (b x y))
  go (Use nam val bod) = Use nam (go val) (\x -> go (bod (Sub x)))
  go (Met uid spn)     = Met uid (map go spn)
  go (Log msg nxt)     = Log (go msg) (go nxt)
  go (Hol nam ctx)     = Hol nam (map go ctx)
  go Set               = Set
  go U64               = U64
  go F64               = F64
  go (Num n)           = Num n
  go (Flt n)           = Flt n
  go (Op2 opr fst snd) = Op2 opr (go fst) (go snd)
  go (Txt txt)         = Txt txt
  go (Lst lst)         = Lst (map go lst)
  go (Nat val)         = Nat val
  go (Var nam idx)     = if lvl == idx then neo else Var nam idx
  go (Src src val)     = Src src (go val)
  go (Sub val)         = val
  goCtr (Ctr nm tele)  = Ctr nm (goTele tele)
  goCse (cnam, cbod)   = (cnam, go cbod)
  goTele (TRet term)   = TRet (go term)
  goTele (TExt k t b)  = TExt k (go t) (\x -> goTele (b x))

-- Replaces a term by another
replace :: Term -> Term -> Term -> Int -> Term
replace old neo term dep = if same old term dep then neo else go term where
  go (All nam inp bod)  = All nam (replace old neo inp dep) (\x -> replace old neo (bod (Sub x)) (dep+1))
  go (Lam nam bod)      = Lam nam (\x -> replace old neo (bod (Sub x)) (dep+1))
  go (App fun arg)      = App (replace old neo fun dep) (replace old neo arg dep)
  go (Ann chk val typ)  = Ann chk (replace old neo val dep) (replace old neo typ dep)
  go (Slf nam typ bod)  = Slf nam (replace old neo typ dep) (\x -> replace old neo (bod (Sub x)) (dep+1))
  go (Ins val)          = Ins (replace old neo val dep)
  go (ADT scp cts typ)  = ADT (map (\x -> replace old neo x (dep+1)) scp) (map goCtr cts) (replace old neo typ dep)
  go (Con nam arg)      = Con nam (map (\(f, t) -> (f, replace old neo t dep)) arg)
  go (Mat cse)          = Mat (map goCse cse)
  go (Swi zer suc)      = Swi (replace old neo zer dep) (replace old neo suc dep)
  go (Map typ)          = Map (replace old neo typ dep)
  go (KVs map def)      = KVs (IM.map (\x -> replace old neo x dep) map) (replace old neo def dep)
  go (Get g n m k b)    = Get g n (replace old neo m dep) (replace old neo k dep) (\x y -> replace old neo (b x y) (dep+2))
  go (Put g n m k v b)  = Put g n (replace old neo m dep) (replace old neo k dep) (replace old neo v dep) (\x y -> replace old neo (b x y) (dep+2))
  go (Ref nam)          = Ref nam
  go (Let nam val bod)  = Let nam (replace old neo val dep) (\x -> replace old neo (bod (Sub x)) (dep+1))
  go (Use nam val bod)  = Use nam (replace old neo val dep) (\x -> replace old neo (bod (Sub x)) (dep+1))
  go (Met uid spn)      = Met uid (map (\x -> replace old neo x (dep+1)) spn)
  go (Log msg nxt)      = Log (replace old neo msg dep) (replace old neo nxt dep)
  go (Hol nam ctx)      = Hol nam (map (\x -> replace old neo x (dep+1)) ctx)
  go Set                = Set
  go U64                = U64
  go F64                = F64
  go (Num n)            = Num n
  go (Flt n)            = Flt n
  go (Op2 opr fst snd)  = Op2 opr (replace old neo fst dep) (replace old neo snd dep)
  go (Txt txt)          = Txt txt
  go (Lst lst)          = Lst (map (\x -> replace old neo x dep) lst)
  go (Nat val)          = Nat val
  go (Var nam idx)      = Var nam idx
  go (Src src val)      = Src src (replace old neo val dep)
  go (Sub val)          = val
  goCtr (Ctr nm tele)   = Ctr nm (goTele tele dep)
  goCse (cnam, cbod)    = (cnam, replace old neo cbod dep)
  goTele (TRet term)  d = TRet (replace old neo term d)
  goTele (TExt k t b) d = TExt k (replace old neo t d) (\x -> goTele (b x) (d+1))

-- Returns true when two terms can definitely never be made identical.
-- TODO: to implement this, just recurse pairwise on the Con constructor,
-- until a different name is found. All other terms are considered compatible.
incompatible :: Term -> Term -> Int -> Bool
incompatible (Con aNam aArg) (Con bNam bArg) dep | aNam /= bNam = True
incompatible (Con aNam aArg) (Con bNam bArg) dep | otherwise    = length aArg == length bArg && any (\(a,b) -> incompatible a b dep) (zip (map snd aArg) (map snd bArg))
incompatible (Src aSrc aVal) b               dep                = incompatible aVal b dep
incompatible a               (Src bSrc bVal) dep                = incompatible a bVal dep
incompatible _               _               _                  = False


================================================
FILE: src/Kind/Parse.hs
================================================
-- //./Type.hs//

module Kind.Parse where

import Data.Char (ord)
import Data.Functor.Identity (Identity)
import Data.List (intercalate, isPrefixOf, uncons, unsnoc, find, transpose)
import Data.Maybe (catMaybes, fromJust, isJust)
import Data.Set (toList, fromList)
import Data.Word
import Debug.Trace
import Highlight (highlightError, highlight)
import Kind.Equal
import Kind.Reduce
import Kind.Show
import Kind.Type
import Prelude hiding (EQ, LT, GT)
import System.Console.ANSI
import Text.Parsec ((<?>), (<|>), getPosition, sourceLine, sourceColumn, getState, setState)
import Text.Parsec.Error (errorPos, errorMessages, showErrorMessages, ParseError, errorMessages, Message(..))
import qualified Control.Applicative as A
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Text.Parsec as P

type Uses     = [(String, String)]
type PState   = (String, Int, Uses)
type Parser a = P.ParsecT String PState Identity a
-- Types used for flattening pattern-matching equations
type Rule     = ([Pattern], Term)
data Pattern  = PVar String | PCtr (Maybe String) String [Pattern] | PNum Word64 | PSuc Word64 String

-- Helper functions that consume trailing whitespace
skip :: Parser ()
skip = P.skipMany (parseSpace <|> parseComment)
  where
    parseSpace = (P.try $ do
      P.space
      return ()) <?> "Space"
    parseComment = (P.try $ do
      P.string "//"
      P.skipMany (P.noneOf "\n")
      P.char '\n'
      return ()) <?> "Comment"

char :: Char -> Parser Char
char c = P.char c

string :: String -> Parser String
string s = P.string s

char_skp :: Char -> Parser Char
char_skp c = P.char c <* skip

string_skp :: String -> Parser String
string_skp s = P.string s <* skip

name_init :: Parser Char
name_init = P.satisfy (`elem` "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz/_.-$")

name_char :: Parser Char
name_char = P.satisfy (`elem` "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789/_.-$")

name :: Parser String
name = (do
  head <- name_init
  tail <- P.many name_char
  return (head : tail)) <?> "Name"

name_skp :: Parser String
name_skp = name <* skip

digit :: Parser Char
digit = P.digit

numeric :: Parser String
numeric = (do
  head <- P.satisfy (`elem` "0123456789")
  tail <- P.many (P.satisfy (`elem` "bx0123456789abcdefABCDEF_"))
  return $ show (read (filter (/= '_') (head : tail)) :: Word64)) <?> "Number"

numeric_skp :: Parser String
numeric_skp = numeric <* skip

oneOf :: String -> Parser Char
oneOf s = P.oneOf s

noneOf :: String -> Parser Char
noneOf s = P.noneOf s

guardChoice :: [(Parser a, Parser ())] -> Parser a -> Parser a
guardChoice []          df = df
guardChoice ((p, g):ps) df = do
  guard <- P.lookAhead $ P.optionMaybe $ P.try g
  case guard of
    Just () -> p
    Nothing -> guardChoice ps df

discard :: Parser a -> Parser ()
discard p = p >> return ()

-- Main parsing functions
doParseTerm :: String -> String -> IO Term
doParseTerm filename input =
  case P.runParser (parseTerm <* P.eof) (filename, 0, []) filename input of
    Left err -> do
      showParseError filename input err
      return $ Ref "bad-parse"
    Right term -> return $ bind (genMetas term) []

doParseUses :: String -> String -> IO Uses
doParseUses filename input =
  case P.runParser (parseUses <* P.eof) (filename, 0, []) filename input of
    Left err -> do
      showParseError filename input err
      return $ []
    Right uses -> return uses

doParseBook :: String -> String -> IO Book
doParseBook filename input = do
  let parser = do
        skip
        uses <- parseUses
        setState (filename, 0, uses)
        parseBook <* P.eof
  case P.runParser parser (filename, 0, []) filename input of
    Left err -> do
      showParseError filename input err
      return M.empty
    Right book -> return book

-- Error handling
extractExpectedTokens :: ParseError -> String
extractExpectedTokens err =
    let expectedMsgs = [msg | Expect msg <- errorMessages err, msg /= "Space", msg /= "Comment"]
    in intercalate " | " expectedMsgs

showParseError :: String -> String -> P.ParseError -> IO ()
showParseError filename input err = do
  let pos = errorPos err
  let lin = sourceLine pos
  let col = sourceColumn pos
  let errorMsg = extractExpectedTokens err
  putStrLn $ setSGRCode [SetConsoleIntensity BoldIntensity] ++ "\nPARSE_ERROR" ++ setSGRCode [Reset]
  putStrLn $ "- expected: " ++ errorMsg
  putStrLn $ "- detected:"
  putStrLn $ highlightError (lin, col) (lin, col + 1) input
  putStrLn $ setSGRCode [SetUnderlining SingleUnderline] ++ filename ++
             setSGRCode [Reset] ++ " " ++ show lin ++ ":" ++ show col

-- Parsing helpers
-- FIXME: currently, this will include suffix trivia. how can we avoid that?
withSrc :: Parser Term -> Parser Term
withSrc parser = do
  ini <- getPosition
  val <- parser
  end <- getPosition
  (nam, _, _) <- P.getState
  skip
  let iniLoc = Loc nam (sourceLine ini) (sourceColumn ini)
  let endLoc = Loc nam (sourceLine end) (sourceColumn end)
  return $ Src (Cod iniLoc endLoc) val

-- Term Parser
-- -----------

-- Main term parser
parseTerm :: Parser Term
parseTerm = do
  skip
  term <- guardChoice
    [ (parseAll,             discard $ string_skp "∀")
    , (parseSwi,             discard $ string_skp "λ" >> string_skp "{" >> string_skp "0")
    , (parseMat,             discard $ string_skp "λ" >> string_skp "{" >> string_skp "#")
    , (parseLam,             discard $ string_skp "λ")
    , (parseEra,             discard $ string_skp "λ")
    , (parseOp2,             discard $ string_skp "(" >> parseOper)
    , (parseMap,             discard $ string_skp "(Map ")
    , (parseApp,             discard $ string_skp "(")
    , (parseSlf,             discard $ string_skp "$(")
    , (parseIns,             discard $ string_skp "~")
    , (parseADT,             discard $ string_skp "#[" <|> string_skp "data[")
    , (parseNat,             discard $ string_skp "#" >> digit)
    , (parseCon,             discard $ string_skp "#" >> name)
    , ((parseUse parseTerm), discard $ string_skp "use ")
    , ((parseLet parseTerm), discard $ string_skp "let ")
    , ((parseGet parseTerm), discard $ string_skp "get ")
    , ((parsePut parseTerm), discard $ string_skp "put ")
    , (parseIf,              discard $ string_skp "if ")
    , (parseWhen,            discard $ string_skp "when ")
    , (parseMatInl,          discard $ string_skp "match ")
    , (parseSwiInl,          discard $ string_skp "switch ")
    , (parseKVs,             discard $ string_skp "{")
    , (parseDo,              discard $ string_skp "do ")
    , (parseSet,             discard $ string_skp "*")
    , (parseFloat,           discard $ string_skp "-" <|> (P.many1 digit >> string_skp "."))
    , (parseNum,             discard $ numeric)
    , (parseTxt,             discard $ string_skp "\"")
    , (parseLst,             discard $ string_skp "[")
    , (parseChr,             discard $ string_skp "'")
    , (parseHol,             discard $ string_skp "?")
    , ((parseLog parseTerm), discard $ string_skp "log ")
    , (parseRef,             discard $ name)
    ] $ fail "Term"
  skip
  parseSuffix term

-- Individual term parsers
parseAll = withSrc $ do
  string_skp "∀"
  era <- P.optionMaybe (char_skp '-')
  char_skp '('
  nam <- name_skp
  char_skp ':'
  inp <- parseTerm
  char_skp ')'
  bod <- parseTerm
  return $ All nam inp (\x -> bod)

parseLam = withSrc $ do
  string_skp "λ"
  era <- P.optionMaybe (char_skp '-')
  nam <- name_skp
  bod <- parseTerm
  return $ Lam nam (\x -> bod)

parseEra = withSrc $ do
  string_skp "λ"
  era <- P.optionMaybe (char_skp '-')
  nam <- char_skp '_'
  bod <- parseTerm
  return $ Lam "_" (\x -> bod)

parseApp = withSrc $ do
  char_skp '('
  fun  <- parseTerm
  args <- P.many $ do
    P.notFollowedBy (char ')')
    era <- P.optionMaybe (char_skp '-')
    arg <- parseTerm
    return (era, arg)
  char ')'
  return $ foldl (\f (era, a) -> App f a) fun args

parseSlf = withSrc $ do
  string_skp "$("
  nam <- name_skp
  char_skp ':'
  typ <- parseTerm
  char_skp ')'
  bod <- parseTerm
  return $ Slf nam typ (\x -> bod)

parseIns = withSrc $ do
  char_skp '~'
  val <- parseTerm
  return $ Ins val

parseADT = withSrc $ do
  P.choice [string_skp "#[", string_skp "data["]
  scp <- P.many parseTerm
  char_skp ']'
  char_skp '{'
  cts <- P.many $ P.try parseADTCtr
  char '}'
  typ <- do
    skip
    char_skp ':'
    parseTerm
  return $ ADT scp cts typ

parseADTCtr :: Parser Ctr
parseADTCtr = do
  char_skp '#'
  name <- name_skp
  tele <- parseTele
  return $ Ctr name tele

parseTele :: Parser Tele
parseTele = do
  fields <- P.option [] $ do
    char_skp '{'
    fields <- P.many $ P.try $ do
      nam <- name_skp
      char_skp ':'
      typ <- parseTerm
      return (nam, typ)
    char_skp '}'
    return fields
  ret <- P.choice
    [ do
        P.try $ char_skp ':'
        parseTerm
    , do
        return (Met 0 [])
    ]
  return $ foldr (\(nam, typ) acc -> TExt nam typ (\x -> acc)) (TRet ret) fields

parseCon = withSrc $ do
  char_skp '#'
  nam <- name
  args <- P.option [] $ P.try $ do
    skip
    char_skp '{'
    args <- P.many $ do
      P.notFollowedBy (char_skp '}')
      name <- P.optionMaybe $ P.try $ do
        name <- name_skp
        char_skp ':'
        return name
      term <- parseTerm
      return (name, term)
    char '}'
    return args
  return $ Con nam args

parseMatCases :: Parser [(String, Term)]
parseMatCases = do
  cse <- P.many $ P.try $ do
    string_skp "#"
    cnam <- name_skp
    args <- P.option [] $ P.try $ do
      char_skp '{'
      names <- P.many name_skp
      char_skp '}'
      return names
    char_skp ':'
    cbod <- parseTerm
    return (cnam, foldr (\arg acc -> Lam arg (\_ -> acc)) cbod args)
  dflt <- P.optionMaybe $ do
    dnam <- P.try $ do
      dnam <- name_skp
      string_skp ":"
      return dnam
    dbod <- parseTerm
    return (dnam, dbod)
  return $ case dflt of
    Just (dnam, dbod) -> cse ++ [("_", (Lam dnam (\_ -> dbod)))]
    Nothing           -> cse

parseSwiCases :: Parser Term
parseSwiCases = do
  cse <- P.many $ P.try $ do
    cnam <- numeric_skp
    char_skp ':'
    cbod <- parseTerm
    return (cnam, cbod)
  dflt <- P.optionMaybe $ do
    dnam <- P.try $ do
      cnam <- numeric_skp
      char_skp '+'
      dnam <- name_skp
      string_skp ":"
      return dnam
    dbod <- parseTerm
    return (dnam, dbod)
  case dflt of
    Just (dnam, dbod) -> return $ build (cse ++ [("_", (Lam dnam (\_ -> dbod)))]) 0
    Nothing           -> return $ build cse 0
  where build :: [(String, Term)] -> Int -> Term
        build []           i               = error "Switch must have at least one case."
        build (("_",t):cs) i               = t
        build ((n,t):cs)   i | read n == i = Swi t (build cs (i+1))
        build ((n,t):cs)   i | otherwise   = error "Switch cases must be in ascending order starting from 0."

parseSwiElim :: Parser Term
parseSwiElim = do
  cases <- parseSwiCases
  return cases

parseSwi = withSrc $ do
  char_skp 'λ'
  char_skp '{'
  P.lookAhead $ P.try $ char_skp '0'
  elim <- parseSwiElim
  char '}'
  return $ elim

parseMat = withSrc $ do
  char_skp 'λ'
  char_skp '{'
  cse <- parseMatCases
  char '}'
  return $ Mat cse

-- TODO: implement the Map parsers
parseMap = withSrc $ do
  string_skp "(Map "
  typ <- parseTerm
  char ')'
  return $ Map typ

parseKVs = withSrc $ do
  char_skp '{'
  kvs <- P.many parseKV
  char_skp '|'
  dft <- parseTerm
  char '}'
  return $ KVs (IM.fromList kvs) dft
  where
    parseKV = do
      key <- read <$> numeric_skp
      char_skp ':'
      val <- parseTerm
      return (key, val)

parseGet parseBody = withSrc $ do
  string_skp "get "
  got <- name_skp
  string_skp "="
  nam <- name_skp
  map <- P.option (Ref nam) $ P.try $ char_skp '@' >> parseTerm
  char_skp '['
  key <- parseTerm
  char_skp ']'
  bod <- parseBody
  return $ Get got nam map key (\x y -> bod)

parsePut parseBody = withSrc $ do
  string_skp "put "
  got <- P.option "_" $ P.try $ do
    got <- name_skp
    string_skp "="
    return got
  nam <- name_skp
  map <- P.option (Ref nam) $ P.try $ char_skp '@' >> parseTerm
  char_skp '['
  key <- parseTerm
  char_skp ']'
  string_skp ":="
  val <- parseTerm
  bod <- parseBody
  return $ Put got nam map key val (\x y -> bod)

parseRef = withSrc $ do
  name <- name
  (_, _, uses) <- P.getState
  let name' = expandUses uses name
  return $ case name' of
    "U64" -> U64
    "F64" -> F64
    "Set" -> Set
    "_"   -> Met 0 []
    _     -> Ref name'

parseLocal :: String -> (String -> Term -> (Term -> Term) -> Term) -> Parser Term -> Parser Term
parseLocal header ctor parseBody = withSrc $ P.choice
  [ parseLocalMch header ctor parseBody
  , parseLocalPar header ctor parseBody
  , parseLocalVal header ctor parseBody
  ]

parseLocalMch :: String -> (String -> Term -> (Term -> Term) -> Term) -> Parser Term -> Parser Term
parseLocalMch header ctor parseBody = do
  P.try $ string_skp (header ++ " #")
  cnam <- name_skp
  char_skp '{'
  args <- P.many name_skp
  char_skp '}'
  char_skp '='
  val <- parseTerm
  bod <- parseBody
  return $ ctor "got" val (\got ->
    App (Mat [(cnam, foldr (\arg acc -> Lam arg (\_ -> acc)) bod args)]) got)

parseLocalPar :: String -> (String -> Term -> (Term -> Term) -> Term) -> Parser Term -> Parser Term
parseLocalPar header ctor parseBody = do
  P.try $ string_skp (header ++ " (")
  head <- name_skp
  tail <- P.many $ do
    char_skp ','
    name_skp
  char_skp ')'
  let (init, last) = maybe ([], head) id $ unsnoc (head : tail)
  char_skp '='
  val <- parseTerm
  bod <- parseBody
  return $ ctor "got" val (\got ->
    App (foldr (\x acc -> Mat [("Pair", Lam x (\_ -> acc))]) (Lam last (\_ -> bod)) init) got)

parseLocalVal :: String -> (String -> Term -> (Term -> Term) -> Term) -> Parser Term -> Parser Term
parseLocalVal header ctor parseBody = do
  P.try $ string_skp (header ++ " ")
  nam <- name_skp
  char_skp '='
  val <- parseTerm
  bod <- parseBody
  return $ ctor nam val (\x -> bod)

parseLet :: Parser Term -> Parser Term
parseLet = parseLocal "let" Let

parseUse :: Parser Term -> Parser Term
parseUse = parseLocal "use" Use

parseSet = withSrc $ char '*' >> return Set

parseFloat = withSrc $ P.try $ do
  -- Parse optional negative sign
  sign <- P.option id $ P.char '-' >> return negate

  -- Parse integer part
  intPart <- P.many1 digit

  -- Parse decimal part (this must succeed, or we fail the whole parser)
  decPart <- do
    char '.'
    P.many1 digit

  -- Parse optional exponent
  expPart <- P.option 0 $ P.try $ do
    oneOf "eE"
    expSign <- P.option '+' (oneOf "+-")
    exp <- read <$> P.many1 digit
    return $ if expSign == '-' then -exp else exp

  -- Combine parts into final float
  let floatStr = intPart ++ "." ++ decPart
  let value = (read floatStr :: Double) * (10 ^^ expPart)

  -- Apply the sign to the final value
  return $ Flt (sign value)

parseNum = withSrc $ do
  val <- numeric
  return $ Num (read (filter (/= '_') val))

parseOp2 = withSrc $ do
  char_skp '('
  opr <- parseOper
  fst <- parseTerm
  snd <- parseTerm
  char ')'
  return $ Op2 opr fst snd

parseLst = withSrc $ do
  char_skp '['
  elems <- P.many parseTerm
  char ']'
  return $ Lst elems

parseTxtChr :: Parser Char
parseTxtChr = P.choice
  [ P.try $ do
      char_skp '\\'
      c <- oneOf "\\\"nrtbf0/\'"
      return $ case c of
        '\\' -> '\\'
        '/'  -> '/'
        '"'  -> '"'
        '\'' -> '\''
        'n'  -> '\n'
        'r'  -> '\r'
        't'  -> '\t'
        'b'  -> '\b'
        'f'  -> '\f'
        '0'  -> '\0'
  , P.try $ do
      string_skp "\\u"
      code <- P.count 4 P.hexDigit
      return $ toEnum (read ("0x" ++ code) :: Int)
  , noneOf "\"\\"
  ]

parseTxt = withSrc $ do
  char '"'
  txt <- P.many parseTxtChr
  char '"'
  return $ Txt txt

parseChr = withSrc $ do
  char '\''
  chr <- parseTxtChr
  char '\''
  return $ Num (fromIntegral $ ord chr)

parseHol = withSrc $ do
  char_skp '?'
  nam <- name_skp
  ctx <- P.option [] $ do
    char_skp '['
    terms <- P.sepBy parseTerm (char_skp ',')
    char ']'
    return terms
  return $ Hol nam ctx

parseLog parseBody = withSrc $ do
  string_skp "log "
  msg <- parseTerm
  val <- parseBody
  return $ Log msg val

parseOper = P.choice
  [ P.try (string_skp "+") >> return ADD
  , P.try (string_skp "-") >> return SUB
  , P.try (string_skp "*") >> return MUL
  , P.try (string_skp "/") >> return DIV
  , P.try (string_skp "%") >> return MOD
  , P.try (string_skp "<<") >> return LSH
  , P.try (string_skp ">>") >> return RSH
  , P.try (string_skp "<=") >> return LTE
  , P.try (string_skp ">=") >> return GTE
  , P.try (string_skp "<") >> return LT
  , P.try (string_skp ">") >> return GT
  , P.try (string_skp "==") >> return EQ
  , P.try (string_skp "!=") >> return NE
  , P.try (string_skp "&") >> return AND
  , P.try (string_skp "|") >> return OR
  , P.try (string_skp "^") >> return XOR
  ] <?> "Binary operator"

parseSuffix :: Term -> Parser Term
parseSuffix term = guardChoice
  [ (parseSuffArr term, discard $ string_skp "->")
  , (parseSuffAnn term, discard $ string_skp "::")
  , (parseSuffEql term, discard $ string_skp "==")
  , (parseSuffPAR term, discard $ string_skp "&")
  , (parseSuffPar term, discard $ string_skp ",")
  , (parseSuffCns term, discard $ string_skp ";;")
  ] $ parseSuffVal term

parseSuffArr :: Term -> Parser Term
parseSuffArr term = do
  P.try $ string_skp "->"
  ret <- parseTerm
  return $ All "_" term (\_ -> ret)

parseSuffAnn :: Term -> Parser Term
parseSuffAnn term = do
  P.try $ string_skp "::"
  typ <- parseTerm
  return $ Ann True term typ

parseSuffEql :: Term -> Parser Term
parseSuffEql term = do
  P.try $ string_skp "=="
  other <- parseTerm
  return $ App (App (App (Ref "Equal") (Met 0 [])) term) other

parseSuffPAR :: Term -> Parser Term
parseSuffPAR fst = do
  P.try $ string_skp "&"
  snd <- parseTerm
  return $ App (App (Ref "Pair") fst) snd

parseSuffPar :: Term -> Parser Term
parseSuffPar fst = do
  P.try $ string_skp ","
  snd <- parseTerm
  return $ Con "Pair" [(Nothing, fst), (Nothing, snd)]

parseSuffCns :: Term -> Parser Term
parseSuffCns head = do
  P.try $ string_skp ";;"
  tail <- parseTerm
  return $ Con "Cons" [(Nothing, head), (Nothing, tail)]

parseSuffVal :: Term -> Parser Term
parseSuffVal term = return term

-- Book Parser
-- -----------

parseBook :: Parser Book
parseBook = M.fromList <$> P.many parseDef

parseDef :: Parser (String, Term)
parseDef = guardChoice
  [ (parseDefADT, discard $ string_skp "data ")
  , (parseDefFun, discard $ string_skp "#" <|> name_skp)
  ] $ fail "Top-level definition"

parseDefADT :: Parser (String, Term)
parseDefADT = do
  (_, _, uses) <- P.getState
  P.try $ string_skp "data "
  name <- name_skp
  let nameA = expandUses uses name
  params <- P.many $ do
    P.try $ char_skp '('
    pname <- name_skp
    char_skp ':'
    ptype <- parseTerm
    char_skp ')'
    return (pname, ptype)
  indices <- P.choice
    [ do
        P.try $ char_skp '~'
        P.many $ do
          P.notFollowedBy (char '{')
          char_skp '('
          iname <- name_skp
          char_skp ':'
          itype <- parseTerm
          char_skp ')'
          return (iname, itype)
    , return []
    ]
  char_skp '{'
  ctrs <- P.many $ P.try parseADTCtr
  char_skp '}'
  let paramTypes = map snd params
  let indexTypes = map snd indices
  let paramNames = map fst params
  let indexNames = map fst indices
  let allParams  = params ++ indices
  let selfType   = foldl (\ acc arg -> App acc (Ref arg)) (Ref nameA) (paramNames ++ indexNames)
  let typeBody   = foldr (\ (pname, ptype) acc -> All pname ptype (\_ -> acc)) Set allParams
  let newCtrs    = map (fillCtrRet selfType) ctrs -- fill ctr type when omitted
  let dataBody   = ADT (map (\ (iNam,iTyp) -> Ref iNam) indices) newCtrs selfType
  let fullBody   = foldr (\ (pname, _) acc -> Lam pname (\_ -> acc)) dataBody allParams
  let term       = bind (genMetas (Ann False fullBody typeBody)) []
  return $
    -- trace ("parsed " ++ nameA ++ " = " ++ (showTermGo False term 0))
    (nameA, term)
  where fillCtrRet  ret (Ctr nm tele)    = Ctr nm (fillTeleRet ret tele)
        fillTeleRet ret (TRet (Met _ _)) = TRet ret
        fillTeleRet _   (TRet ret)       = TRet ret
        fillTeleRet ret (TExt nm tm bod) = TExt nm tm (\x -> fillTeleRet ret (bod x)) -- FIXME: 'bod x'?

parseDefFun :: Parser (String, Term)
parseDefFun = do
  numb <- P.optionMaybe $ char_skp '#'
  name <- name_skp
  typ <- P.optionMaybe $ do
    char_skp ':'
    t <- parseTerm
    return t
  val <- guardChoice
    [ (parseDefFunSingle, discard $ char_skp '=')
    , (parseDefFunRules,  discard $ char_skp '|')
    ] parseDefFunTest
  (filename, count, uses) <- P.getState
  let name0 = expandUses uses name
  let name1 = if isJust numb then name0 ++ "#" ++ show count else name0
  P.setState (filename, if isJust numb then count + 1 else count, uses)
  case typ of
    Nothing -> return (name1, bind (genMetas val) [])
    Just t  -> return (name1, bind (genMetas (Ann False val t)) [])

parseDefFunSingle :: Parser Term
parseDefFunSingle = do
  char_skp '='
  val <- parseTerm
  return val

parseDefFunRules :: Parser Term
parseDefFunRules = withSrc $ do
  rules <- P.many1 (parseRule 0)
  let flat = flattenDef rules 0
  return
    -- $ trace ("DONE: " ++ showTerm flat)
    flat

parseDefFunTest :: Parser Term
parseDefFunTest = return (Con "Refl" [])

parseRule :: Int -> Parser Rule
parseRule dep = do
  P.try $ do
    P.count dep $ char_skp '.'
    char_skp '|'
  pats <- P.many parsePattern
  body <- P.choice 
    [ withSrc $ P.try $ do
      string_skp "with "
      wth <- P.many1 $ P.notFollowedBy (char_skp '.') >> parseTerm
      rul <- P.many1 $ parseRule (dep + 1)
      return $ flattenWith dep wth rul
    , P.try $ do
      char_skp '='
      body <- parseTerm
      return body
    ]
  return $ (pats, body)

parsePattern :: Parser Pattern
parsePattern = do
  P.notFollowedBy $ string_skp "with "
  pat <- guardChoice
    [ (parsePatPrn, discard $ string_skp "(")
    , (parsePatNat, discard $ string_skp "#" >> numeric_skp)
    , (parsePatLst, discard $ string_skp "[")
    , (parsePatCon, discard $ string_skp "#" <|> (name_skp >> string_skp "@"))
    , (parsePatTxt, discard $ string_skp "\"")
    , (parsePatSuc, discard $ numeric_skp >> char_skp '+')
    , (parsePatNum, discard $ numeric_skp)
    , (parsePatVar, discard $ name_skp)
    ] $ fail "Pattern-matching"
  parsePatSuffix pat

parsePatSuffix :: Pattern -> Parser Pattern
parsePatSuffix pat = P.choice
  [ parsePatSuffPar pat
  , parsePatSuffCns pat
  , return pat
  ]

parsePatSuffPar :: Pattern -> Parser Pattern
parsePatSuffPar fst = do
  P.try $ string_skp ","
  snd <- parsePattern
  return $ PCtr Nothing "Pair" [fst, snd]

parsePatSuffCns :: Pattern -> Parser Pattern
parsePatSuffCns head = do
  P.try $ string_skp ";;"
  tail <- parsePattern
  return $ PCtr Nothing "Cons" [head, tail]

parsePatPrn :: Parser Pattern
parsePatPrn = do
  string_skp "("
  pat <- parsePattern
  string_skp ")"
  return pat

parsePatNat :: Parser Pattern
parsePatNat = do
  char_skp '#'
  num <- numeric_skp
  let n = read num
  return $ (foldr (\_ acc -> PCtr Nothing "Succ" [acc]) (PCtr Nothing "Zero" []) [1..n])

parsePatLst :: Parser Pattern
parsePatLst = do
  char_skp '['
  elems <- P.many parsePattern
  char_skp ']'
  return $ foldr (\x acc -> PCtr Nothing "Cons" [x, acc]) (PCtr Nothing "Nil" []) elems

parsePatTxt :: Parser Pattern
parsePatTxt = do
  char '"'
  txt <- P.many parseTxtChr
  char '"'
  return $ foldr (\x acc -> PCtr Nothing "Cons" [PNum (toEnum (ord x)), acc]) (PCtr Nothing "Nil" []) txt

parsePatPar :: Parser Pattern
parsePatPar = do
  char_skp '('
  head <- parsePattern
  tail <- P.many $ do
    char_skp ','
    parsePattern
  char_skp ')'
  let (init, last) = maybe ([], head) id (unsnoc (head : tail))
  return $ foldr (\x acc -> PCtr Nothing "Pair" [x, acc]) last init

parsePatCon :: Parser Pattern
parsePatCon = do
  name <- P.optionMaybe $ P.try $ do
    name <- name_skp
    char_skp '@'
    return name
  char_skp '#'
  cnam <- name_skp
  args <- P.option [] $ P.try $ do
    char_skp '{'
    args <- P.many parsePattern
    char_skp '}'
    return args
  return $ (PCtr name cnam args)

parsePatNum :: Parser Pattern
parsePatNum = do
  num <- numeric_skp
  return $ (PNum (read num))

parsePatSuc :: Parser Pattern
parsePatSuc = do
  num <- numeric_skp
  char_skp '+'
  nam <- name_skp
  return $ (PSuc (read num) nam)

parsePatVar :: Parser Pattern
parsePatVar = do
  name <- name_skp
  return $ (PVar name)

parseUses :: Parser Uses
parseUses = P.many $ P.try $ do
  string_skp "use "
  long <- name_skp
  string_skp "as "
  short <- name_skp
  return (short, long)

expandUses :: Uses -> String -> String
expandUses ((short, long):uses) name
  | short == name                    = long
  | (short ++ "/") `isPrefixOf` name = long ++ drop (length short) name
  | otherwise                        = expandUses uses name
expandUses [] name                   = name

-- Syntax Sugars
-- -------------

parseDo :: Parser Term
parseDo = withSrc $ do
  string_skp "do "
  monad <- name_skp
  char_skp '{'
  skip
  (_, _, uses) <- P.getState
  body <- parseStmt (expandUses uses monad)
  char '}'
  return body

parseStmt :: String -> Parser Term
parseStmt monad = guardChoice
  [ (parseDoFor monad,           discard $ string_skp "for " <|> (string_skp "ask" >> name_skp >> string_skp "=" >> string_skp "for"))
  , (parseDoAsk monad,           discard $ string_skp "ask ")
  , (parseDoRet monad,           discard $ string_skp "ret ")
  , (parseLet (parseStmt monad), discard $ string_skp "let ")
  , (parseUse (parseStmt monad), discard $ string_skp "use ")
  , (parseLog (parseStmt monad), discard $ string_skp "log ")
  ] parseTerm

parseDoAsk :: String -> Parser Term
parseDoAsk monad = guardChoice
  [ (parseDoAskMch monad, discard $ string_skp "ask #")
  , (parseDoAskPar monad, discard $ string_skp "ask (" >> name_skp >> string_skp ",")
  , (parseDoAskVal monad, discard $ string_skp "ask ")
  ] $ fail "'ask' statement"

parseDoAskMch :: String -> Parser Term
parseDoAskMch monad = do
  string_skp "ask #"
  cnam <- name_skp
  char_skp '{'
  args <- P.many name_skp
  char_skp '}'
  char_skp '='
  val <- parseTerm
  next <- parseStmt monad
  (_, _, uses) <- P.getState
  return $ App
    (App (App (App (Ref (monad ++ "/bind")) (Met 0 [])) (Met 0 [])) val)
    (Lam "got" (\got ->
      App (Mat [(cnam, foldr (\arg acc -> Lam arg (\_ -> acc)) next args)]) got))

parseDoAskPar :: String -> Parser Term
parseDoAskPar monad = do
  string_skp "ask ("
  head <- name_skp
  tail <- P.many $ do
    char_skp ','
    name_skp
  char_skp ')'
  let (init, last) = maybe ([], head) id $ unsnoc (head : tail)
  char_skp '='
  val <- parseTerm
  next <- parseStmt monad
  (_, _, uses) <- P.getState
  return $ App
    (App (App (App (Ref (monad ++ "/bind")) (Met 0 [])) (Met 0 [])) val)
    (foldr (\x acc -> Mat [("Pair", Lam x (\_ -> acc))]) (Lam last (\_ -> next)) init)

parseDoAskVal :: String -> Parser Term
parseDoAskVal monad = P.choice
  [ parseDoAskValNamed monad
  , parseDoAskValAnon monad
  ]

parseDoAskValNamed :: String -> Parser Term
parseDoAskValNamed monad = P.try $ do
  string_skp "ask "
  nam <- name_skp
  char_skp '='
  exp <- parseTerm
  next <- parseStmt monad
  (_, _, uses) <- P.getState
  return $ App
    (App (App (App (Ref (monad ++ "/bind")) (Met 0 [])) (Met 0 [])) exp)
    (Lam nam (\_ -> next))

parseDoAskValAnon :: String -> Parser Term
parseDoAskValAnon monad = P.try $ do
  string_skp "ask "
  exp <- parseTerm
  next <- parseStmt monad
  (_, _, uses) <- P.getState
  return $ App
    (App (App (App (Ref (monad ++ "/bind")) (Met 0 [])) (Met 0 [])) exp)
    (Lam "_" (\_ -> next))

parseDoRet :: String -> Parser Term
parseDoRet monad = do
  string_skp "ret "
  exp <- parseTerm
  (_, _, uses) <- P.getState
  return $ App (App (Ref (monad ++ "/pure")) (Met 0 [])) exp

parseDoFor :: String -> Parser Term
parseDoFor monad = do
  (stt, nam, lst, loop, body) <- P.choice
    [ do
        stt <- P.try $ do
          string_skp "ask "
          stt <- name_skp
          string_skp "="
          string_skp "for"
          return stt
        nam <- name_skp
        string_skp "in"
        lst <- parseTerm
        char_skp '{'
        loop <- parseStmt monad
        char_skp '}'
        body <- parseStmt monad
        return (Just stt, nam, lst, loop, body)
    , do
        P.try $ string_skp "for "
        nam <- name_skp
        string_skp "in"
        lst <- parseTerm
        char_skp '{'
        loop <- parseStmt monad
        char_skp '}'
        body <- parseStmt monad
        return (Nothing, nam, lst, loop, body) ]
  let f0 = Ref "List/for"
  let f1 = App f0 (Met 0 [])
  let f2 = App f1 (Ref (monad ++ "/Monad"))
  let f3 = App f2 (Met 0 [])
  let f4 = App f3 (Met 0 [])
  let f5 = App f4 lst
  let f6 = App f5 (maybe (Num 0) Ref stt)
  let f7 = App f6 (Lam (maybe "" id stt) (\s -> Lam nam (\_ -> loop)))
  let b0 = Ref (monad ++ "/bind")
  let b1 = App b0 (Met 0 [])
  let b2 = App b1 (Met 0 [])
  let b3 = App b2 f7
  let b4 = App b3 (Lam (maybe "" id stt) (\_ -> body))
  return b4

-- If-Then-Else
-- ------------

-- if cond { t } else { f }
-- --------------------------------- desugars to
-- match cond { #True: t #False: f }

parseIf = withSrc $ do
  string_skp "if "
  cond <- parseTerm
  t <- parseBranch True
  string_skp "else"
  f <- P.choice [parseBranch False, parseIf]
  return $ App (Mat [("True", t), ("False", f)]) cond
  where
    parseBranch isIf = P.choice
      [ do
          string_skp "do "
          monad <- name_skp
          char_skp '{'
          (_, _, uses) <- P.getState
          t <- parseStmt (expandUses uses monad)
          if isIf then char_skp '}' else char '}'
          return t
      , do
          char_skp '{'
          t <- parseTerm
          if isIf then char_skp '}' else char '}'
          return t
      ]

-- When
-- ----

-- when fn x { c0: v0 c1: v1 } else { df }
-- -------------------------------------------------------- desugars to
-- if (fn x c0) { v0 } else if (fn x c1) { v1 } else { df }

parseWhen = withSrc $ do
  string_skp "when "
  fun <- parseTerm
  val <- parseTerm
  char_skp '{'
  cases <- P.many $ do
    cond <- parseTerm
    char_skp ':'
    body <- parseTerm
    return (cond, body)
  char_skp '}'
  string_skp "else"
  char_skp '{'
  elseCase <- parseTerm
  char '}'
  return $ foldr
    (\ (cond, body) acc -> App
      (Mat [("True", body), ("False", acc)])
      (App (App fun val) cond))
    elseCase
    cases

-- Match
-- -----

parseMatInl :: Parser Term
parseMatInl = withSrc $ do
  string_skp "match "
  x <- parseTerm
  char_skp '{'
  cse <- parseMatCases
  char '}'
  return $ App (Mat cse) x

parseSwiInl :: Parser Term
parseSwiInl = withSrc $ do
  string_skp "switch "
  x <- parseTerm
  char_skp '{'
  cse <- parseSwiCases
  char '}'
  return $ App cse x

-- Nat
-- ---

parseNat :: Parser Term
parseNat = withSrc $ do
  char_skp '#'
  num <- P.many1 digit
  return $ Nat (read num)

-- Flattener
-- ---------

-- FIXME: the functions below are still a little bit messy and can be improved

-- Flattener for pattern matching equations
flattenDef :: [Rule] -> Int -> Term
flattenDef rules depth =
  let (pats, bods) = unzip rules
  in flattenRules pats bods depth

flattenWith :: Int -> [Term] -> [Rule] -> Term
flattenWith dep wth rul =
  -- Wrap the 'with' arguments and patterns in Pairs since the type checker only takes one match argument.
  let wthA = foldr1 (\x acc -> Ann True (Con "Pair" [(Nothing, x), (Nothing, acc)]) (App (App (Ref "Pair") (Met 0 [])) (Met 0 []))) wth
      rulA = map (\(pat, wth) -> ([foldr1 (\x acc -> PCtr Nothing "Pair" [x, acc]) pat], wth)) rul
      bod  = flattenDef rulA (dep + 1)
  in App bod wthA

flattenRules :: [[Pattern]] -> [Term] -> Int -> Term
flattenRules ([]:mat)   (bod:bods) depth = bod
flattenRules (pats:mat) (bod:bods) depth
  | all isVar col                 = flattenVarCol col mat' (bod:bods) (depth + 1)
  | not (null (getColCtrs col))   = flattenAdtCol col mat' (bod:bods) (depth + 1)
  | isJust (fst (getColSucc col)) = flattenNumCol col mat' (bod:bods) (depth + 1)
  | otherwise                     = error "invalid pattern matching function"
  where (col,mat') = getCol (pats:mat)
flattenRules _ _ _ = error "internal error"

-- Flattens a column with only variables
flattenVarCol :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> Term
flattenVarCol col mat bods depth =
  -- trace (replicate (depth * 2) ' ' ++ "flattenVarCol: col = " ++ show col ++ ", depth = " ++ show depth) $
  let nam = maybe "_" id (getVarColName col)
      bod = flattenRules mat bods depth
  in Lam nam (\x -> bod)

-- Flattens a column with constructors and possibly variables
flattenAdtCol :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> Term
flattenAdtCol col mat bods depth =
  -- trace (replicate (depth * 2) ' ' ++ "flattenAdtCol: col = " ++ show col ++ ", depth = " ++ show depth) $
  let ctr = map (makeCtrCase col mat bods depth) (getColCtrs col)
      dfl = makeDflCase col mat bods depth
      nam = getMatNam col
  in case nam of
    (Just nam) -> (Lam nam (\x -> App (Mat (ctr++dfl)) (Ref nam)))
    Nothing    -> Mat (ctr++dfl)

-- Creates a constructor case: '#Name: body'
makeCtrCase :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> String -> (String, Term)
makeCtrCase col mat bods depth ctr =
  -- trace (replicate (depth * 2) ' ' ++ "makeCtrCase: col = " ++ show col ++ ", mat = " ++ show mat ++ ", bods = " ++ show (map showTerm bods) ++ ", depth = " ++ show depth ++ ", ctr = " ++ ctr) $
  let var           = getCtrColNames col ctr
      (mat', bods') = foldr (go var) ([], []) (zip3 col mat bods)
      bod           = flattenRules mat' bods' (depth + 1)
  in (ctr, bod)
  where go var ((PCtr nam cnam ps), pats, bod) (mat, bods)
          | cnam == ctr = ((ps ++ pats):mat, bod:bods)
          | otherwise  = (mat, bods)
        go var ((PVar "_"), pats, bod) (mat, bods) =
          let pat = map (maybe (PVar "_") PVar) var
          in ((pat ++ pats):mat, bod:bods)
        go var ((PVar nam), pats, bod) (mat, bods) =
          let vr2 = [maybe (nam++"."++show i) id vr | (vr, i) <- zip var [0..]]
              pat = map PVar vr2
              bo2 = Use nam (Con ctr (map (\x -> (Nothing, Ref x)) vr2)) (\x -> bod)
          in ((pat ++ pats):mat, bo2:bods)
        go var (_, pats, bod) (mat, bods) =
          (mat, bods)

-- Creates a default case: '#_: body'
makeDflCase :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> [(String, Term)]
makeDflCase col mat bods depth =
  -- trace (replicate (depth * 2) ' ' ++ "makeDflCase: col = " ++ show col ++ ", depth = " ++ show depth) $
  let (mat', bods') = foldr go ([], []) (zip3 col mat bods) in
  if null bods' then [] else [("_", flattenRules mat' bods' (depth + 1))]
  where go ((PVar nam), pats, bod) (mat, bods) = (((PVar nam):pats):mat, bod:bods)
        go (_,          pats, bod) (mat, bods) = (mat, bods)

flattenNumCol :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> Term
flattenNumCol col mat bods depth =
  -- Find the succ case with the value
  let (suc, var) = getColSucc col
      sucA       = fromJust suc
      varA       = maybe ("%n-" ++ show sucA) id var
      numCs      = map (makeNumCase col mat bods depth) [0..sucA-1]
      sucCs      = (makeSucCase col mat bods depth sucA varA)
  in foldr (\x acc -> Swi x acc) sucCs numCs

makeNumCase :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> Word64 -> Term
makeNumCase col mat bods depth num =
  let (mat', bods') = foldr go ([], []) (zip3 col mat bods)
  in if null bods' then error $ "missing case for " ++ show num
     else (flattenRules mat' bods' (depth + 1))
  where go ((PNum val), pats, bod) (mat, bods)
          | val == num = (pats:mat, bod:bods)
          | otherwise  = (mat, bods)
        go ((PVar "_"), pats, bod) (mat, bods) =
          (pats:mat, bod:bods)
        go ((PVar nam), pats, bod) (mat, bods) =
          let bod' = Use nam (Num num) (\x -> bod)
          in (pats:mat, bod':bods)
        go (_, pats, bod) (mat, bods) =
          (mat, bods)

makeSucCase :: [Pattern] -> [[Pattern]] -> [Term] -> Int -> Word64 -> String -> Term
makeSucCase col mat bods depth suc var =
  let (mat', bods') = foldr go ([], []) (zip3 col mat bods)
      bod           = if null bods' then error $ "missing case for " ++ show suc ++ "+" ++ var
                      else (flattenRules mat' bods' (depth + 1))
  in Lam var (\x -> bod)
  where go ((PSuc _ _), pats, bod) (mat, bods) = (pats:mat, bod:bods)
        go ((PVar "_"), pats, bod) (mat, bods) = (pats:mat, bod:bods)
        go ((PVar nam), pats, bod) (mat, bods) = 
          let bodA = Use nam (Op2 ADD (Num suc) (Ref var)) (\x -> bod)
          in (pats:mat, bodA:bods)
        go (_, pats, bod)          (mat, bods) = (mat, bods)

-- Helper Functions

isVar :: Pattern -> Bool
isVar (PVar _) = True
isVar _        = False

getCol :: [[Pattern]] -> ([Pattern], [[Pattern]])
getCol (pats:mat) = unzip (catMaybes (map uncons (pats:mat)))

getColCtrs :: [Pattern] -> [String]
getColCtrs col = toList . fromList $ foldr (\pat acc -> case pat of (PCtr _ cnam _) -> cnam:acc ; _ -> acc) [] col

getVarColName :: [Pattern] -> Maybe String
getVarColName col = foldr (A.<|>) Nothing $ map go col
  where go (PVar "_") = Nothing
        go (PVar nam) = Just nam
        go _          = Nothing

-- For a column of patterns that will become a Mat,
-- return the name of the inner fields or Nothing if they are also Mats.
getCtrColNames :: [Pattern] -> String -> [Maybe String]
getCtrColNames col ctr = 
  let mat = foldr go [] col
  in map getVarColName (transpose mat)
  where go (PCtr nam cnam ps) acc
          | cnam == ctr = ps:acc
          | otherwise   = acc
        go _ acc        = acc

getMatNam :: [Pattern] -> Maybe String
getMatNam (PCtr (Just nam) _ _:_) = Just nam
getMatNam (_:col)                 = getMatNam col
getMatNam []                      = Nothing

-- If theres a PSuc, it returns (Just val, Just nam)
-- If there a PNum a PVar but no PSuc, it returns (Just (max val + 1), Nothing)
-- Otherwise, it returns (Nothing, Nothing)
getColSucc :: [Pattern] -> (Maybe Word64, Maybe String)
getColSucc pats =
  case findSuc pats of
    Just (val, name) -> (Just val, Just name)
    Nothing          -> case (maxNum pats Nothing) of
      Just maxVal -> (Just (maxVal + 1), Nothing) 
      Nothing     -> (Nothing, Nothing)
  where
    findSuc []                = Nothing
    findSuc (PSuc val name:_) = Just (val, name)
    findSuc (_:rest)          = findSuc rest

    maxNum []            acc        = acc
    maxNum (PNum val:ps) Nothing    = maxNum ps (Just val)
    maxNum (PNum val:ps) (Just max) = maxNum ps (Just (if val > max then val else max))
    maxNum (_:ps)        acc        = maxNum ps acc


================================================
FILE: src/Kind/Reduce.hs
================================================
-- //./Type.hs//

module Kind.Reduce where

import Prelude hiding (EQ, LT, GT)
import Data.Bits ( (.&.), (.|.), xor, shiftL, shiftR )
import Data.Char (ord)
import Data.Fixed (mod')
import Debug.Trace
import Kind.Show
import Kind.Type

import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as IM

-- for exitting on undefined ref (should be handled better)
import System.Exit (exitWith, ExitCode(ExitFailure))
import System.IO.Unsafe (unsafePerformIO)

-- Evaluation
-- ----------

-- Evaluates a term to weak normal form
-- 'lv' defines when to expand refs: 0 = never, 1 = on redexes
reduce :: Book -> Fill -> Int -> Term -> Term
reduce book fill lv term = red term where

  red (App fun arg)     = app (red fun) arg
  red (Ann chk val typ) = red val
  red (Ins val)         = red val
  red (Ref nam)         = ref nam
  red (Let nam val bod) = red (bod (red val))
  red (Use nam val bod) = red (bod (red val))
  red (Op2 opr fst snd) = op2 opr (red fst) (red snd)
  red (Txt val)         = txt val
  red (Lst val)         = lst val
  red (Nat val)         = nat val
  red (Src src val)     = red val
  red (Met uid spn)     = met uid spn
  red (Log msg nxt)     = log msg nxt
  red (Get g n m k b)   = get g n (red m) (red k) b
  red (Put g n m k v b) = put g n (red m) (red k) v b
  red val               = val

  app (Ref nam)     arg | lv > 0 = app (ref nam) arg
  app (Met uid spn) arg = red (Met uid (spn ++ [arg]))
  app (Lam nam bod) arg = red (bod (reduce book fill 0 arg))
  app (Mat cse)     arg = mat cse (red arg)
  app (Swi zer suc) arg = swi zer suc (red arg)
  app fun           arg = App fun arg

  mat cse (Con cnam carg) = case lookup cnam cse of
    Just cx -> red (foldl App cx (map snd carg))
    Nothing -> case lookup "_" cse of
      Just df -> red (App df (Con cnam carg))
      Nothing -> error $ "Constructor " ++ cnam ++ " not found in pattern match and no default case '_' provided :" ++ (showTermGo True (Mat cse) 0)
  mat cse arg = App (Mat cse) arg

  swi zer suc (Num 0)             = red zer
  swi zer suc (Num n)             = red (App suc (Num (n - 1)))
  swi zer suc (Op2 ADD (Num 1) k) = red (App suc k)
  swi zer suc val                 = App (Swi zer suc) val

  met uid spn = case IM.lookup uid fill of
    Just val -> red (case spn of
      []       -> val
      (x : xs) -> foldl App val spn)
    Nothing  -> Met uid spn

  op2 op  (Ref nam) (Num snd) | lv > 0 = op2 op (ref nam) (Num snd)
  op2 op  (Num fst) (Ref nam) | lv > 0 = op2 op (Num fst) (ref nam)
  op2 ADD (Num fst) (Num snd) = Num (fst + snd)
  op2 SUB (Num fst) (Num snd) = Num (fst - snd)
  op2 MUL (Num fst) (Num snd) = Num (fst * snd)
  op2 DIV (Num fst) (Num snd) = Num (div fst snd)
  op2 MOD (Num fst) (Num snd) = Num (mod fst snd)
  op2 EQ  (Num fst) (Num snd) = Num (if fst == snd then 1 else 0)
  op2 NE  (Num fst) (Num snd) = Num (if fst /= snd then 1 else 0)
  op2 LT  (Num fst) (Num snd) = Num (if fst < snd then 1 else 0)
  op2 GT  (Num fst) (Num snd) = Num (if fst > snd then 1 else 0)
  op2 LTE (Num fst) (Num snd) = Num (if fst <= snd then 1 else 0)
  op2 GTE (Num fst) (Num snd) = Num (if fst >= snd then 1 else 0)
  op2 AND (Num fst) (Num snd) = Num (fst .&. snd)
  op2 OR  (Num fst) (Num snd) = Num (fst .|. snd)
  op2 XOR (Num fst) (Num snd) = Num (fst `xor` snd)
  op2 LSH (Num fst) (Num snd) = Num (shiftL fst (fromIntegral snd))
  op2 RSH (Num fst) (Num snd) = Num (shiftR fst (fromIntegral snd))
  op2 op  (Ref nam) (Flt snd)  | lv > 0 = op2 op (ref nam) (Flt snd)
  op2 op  (Flt fst) (Ref nam)  | lv > 0 = op2 op (Flt fst) (ref nam)
  op2 ADD (Flt fst) (Flt snd) = Flt (fst + snd)
  op2 SUB (Flt fst) (Flt snd) = Flt (fst - snd)
  op2 MUL (Flt fst) (Flt snd) = Flt (fst * snd)
  op2 DIV (Flt fst) (Flt snd) = Flt (fst / snd)
  op2 MOD (Flt fst) (Flt snd) = Flt (mod' fst snd)
  op2 EQ  (Flt fst) (Flt snd) = Num (if fst == snd then 1 else 0)
  op2 NE  (Flt fst) (Flt snd) = Num (if fst /= snd then 1 else 0)
  op2 LT  (Flt fst) (Flt snd) = Num (if fst < snd then 1 else 0)
  op2 GT  (Flt fst) (Flt snd) = Num (if fst > snd then 1 else 0)
  op2 LTE (Flt fst) (Flt snd) = Num (if fst <= snd then 1 else 0)
  op2 GTE (Flt fst) (Flt snd) = Num (if fst >= snd then 1 else 0)
  op2 AND (Flt _)   (Flt _)   = error "Bitwise AND not supported for floating-point numbers"
  op2 OR  (Flt _)   (Flt _)   = error "Bitwise OR not supported for floating-point numbers"
  op2 XOR (Flt _)   (Flt _)   = error "Bitwise XOR not supported for floating-point numbers"
  op2 opr fst       snd       = Op2 opr fst snd

  ref nam | lv > 0 = case M.lookup nam book of
    Just val -> red val
    Nothing  -> Con ("undefined-reference:"++nam) []
  ref nam = Ref nam

  txt []     = red (Con "Nil" [])
  txt (x:xs) = red (Con "Cons" [(Nothing, Num (toEnum (ord x))), (Nothing, Txt xs)])

  lst []     = red (Con "Nil" [])
  lst (x:xs) = red (Con "Cons" [(Nothing, x), (Nothing, Lst xs)])
  
  nat 0 = Con "Zero" []
  nat n = Con "Succ" [(Nothing, Nat (n - 1))]

  log msg nxt = logMsg book fill lv msg msg nxt ""

  get g n (KVs kvs d) (Num k) b = case IM.lookup (fromIntegral k) kvs of
    Just v  -> red (b v (KVs kvs d))
    Nothing -> red (b d (KVs kvs d))
  get g n m k b = Get g n m k b

  put g n (KVs kvs d) (Num k) v b = case IM.lookup (fromIntegral k) kvs of
    Just o  -> red (b o (KVs (IM.insert (fromIntegral k) v kvs) d))
    Nothing -> red (b d (KVs (IM.insert (fromIntegral k) v kvs) d))
  put g n m k v b = Put g n m k v b

-- Logging
-- -------

logMsg :: Book -> Fill -> Int -> Term -> Term -> Term -> String -> Term
logMsg book fill lv msg' msg nxt txt =
  case (reduce book fill 2 msg) of
    Con "Cons" [(_, head), (_, tail)] -> case (reduce book fill lv head) of
      Num chr -> logMsg book fill lv msg' tail nxt (txt ++ [toEnum (fromIntegral chr)])
      _       -> trace (">> " ++ (showTerm (normal book fill 1 msg' 0))) $ (reduce book fill lv nxt)
    Con "Nil" [] ->
      trace txt (reduce book fill lv nxt)
    bad ->
      trace (">> " ++ (showTerm (normal book fill 1 msg' 0))) $ (reduce book fill lv nxt)

-- Normalization
-- -------------

-- Evaluates a term to full normal form
normal :: Book -> Fill -> Int -> Term -> Int -> Term
normal book fill lv term dep = go (reduce book fill lv term) dep where
  go (All nam inp bod) dep =
    let nf_inp = normal book fill lv inp dep in
    let nf_bod = \x -> normal book fill lv (bod (Var nam dep)) (dep + 1) in
    All nam nf_inp nf_bod
  go (Lam nam bod) dep =
    let nf_bod = \x -> normal book fill lv (bod (Var nam dep)) (dep + 1) in
    Lam nam nf_bod
  go (App fun arg) dep =
    let nf_fun = normal book fill lv fun dep in
    let nf_arg = normal book fill lv arg dep in
    App nf_fun nf_arg
  go (Ann chk val typ) dep =
    let nf_val = normal book fill lv val dep in
    let nf_typ = normal book fill lv typ dep in
    Ann chk nf_val nf_typ
  go (Slf nam typ bod) dep =
    let nf_bod = \x -> normal book fill lv (bod (Var nam dep)) (dep + 1) in
    Slf nam typ nf_bod
  go (Ins val) dep =
    let nf_val = normal book fill lv val dep in
    Ins nf_val
  go (ADT scp cts typ) dep =
    let go_ctr = (\ (Ctr nm tele) ->
          let nf_tele = normalTele book fill lv tele dep in
          Ctr nm nf_tele) in
    let nf_scp = map (\x -> normal book fill lv x dep) scp in
    let nf_cts = map go_ctr cts in
    let nf_typ = normal book fill lv typ dep in
    ADT nf_scp nf_cts nf_typ
  go (Con nam arg) dep =
    let nf_arg = map (\(f, t) -> (f, normal book fill lv t dep)) arg in
    Con nam nf_arg
  go (Mat cse) dep =
    let nf_cse = map (\(cnam, cbod) -> (cnam, normal book fill lv cbod dep)) cse in
    Mat nf_cse
  go (Swi zer suc) dep =
    let nf_zer = normal book fill lv zer dep in
    let nf_suc = normal book fill lv suc dep in
    Swi nf_zer nf_suc
  go (Ref nam) dep = Ref nam
  go (Let nam val bod) dep =
    let nf_val = normal book fill lv val dep in
    let nf_bod = \x -> normal book fill lv (bod (Var nam dep)) (dep + 1) in
    Let nam nf_val nf_bod
  go (Use nam val bod) dep =
    let nf_val = normal book fill lv val dep in
    let nf_bod = \x -> normal book fill lv (bod (Var nam dep)) (dep + 1) in
    Use nam nf_val nf_bod
  go (Hol nam ctx) dep = Hol nam ctx
  go Set dep = Set
  go U64 dep = U64
  go F64 dep = F64
  go (Num val) dep = Num val
  go (Flt val) dep = Flt val
  go (Op2 opr fst snd) dep =
    let nf_fst = normal book fill lv fst dep in
    let nf_snd = normal book fill lv snd dep in
    Op2 opr nf_fst nf_snd
  go (Map typ) dep =
    let nf_typ = normal book fill lv typ dep in
    Map nf_typ
  go (KVs kvs def) dep =
    let nf_kvs = IM.map (\x -> normal book fill lv x dep) kvs in
    let nf_def = normal book fill lv def dep in
    KVs nf_kvs nf_def
  go (Get g n m k b) dep =
    let nf_m = normal book fill lv m dep in
    let nf_k = normal book fill lv k dep in
    let nf_b = \v s -> normal book fill lv (b (Var g dep) (Var n dep)) (dep + 2) in
    Get g n nf_m nf_k nf_b
  go (Put g n m k v b) dep =
    let nf_m = normal book fill lv m dep in
    let nf_k = normal book fill lv k dep in
    let nf_v = normal book fill lv v dep in
    let nf_b = \o s -> normal book fill lv (b (Var g dep) (Var n dep)) (dep + 2) in
    Put g n nf_m nf_k nf_v nf_b
  go (Txt val) dep = Txt val
  go (Lst val) dep =
    let nf_val = map (\x -> normal book fill lv x dep) val in
    Lst nf_val
  go (Nat val) dep = Nat val
  go (Var nam idx) dep = Var nam idx
  go (Src src val) dep =
    let nf_val = normal book fill lv val dep in
    Src src nf_val
  go (Met uid spn) dep = Met uid spn -- TODO: normalize spine
  go (Log msg nxt) dep =
    let nf_msg = normal book fill lv msg dep in
    let nf_nxt = normal book fill lv nxt dep in
    Log nf_msg nf_nxt

normalTele :: Book -> Fill -> Int -> Tele -> Int -> Tele
normalTele book fill lv tele dep = case tele of
  TRet term ->
    let nf_term = normal book fill lv term dep in
    TRet nf_term
  TExt nam typ bod ->
    let nf_typ = normal book fill lv typ dep in
    let nf_bod = \x -> normalTele book fill lv (bod (Var nam dep)) (dep + 1) in
    TExt nam nf_typ nf_bod

-- Binding
-- -------

-- Binds quoted variables to bound HOAS variables
bind :: Term -> [(String,Term)] -> Term
bind (All nam inp bod) ctx =
  let inp' = bind inp ctx in
  let bod' = \x -> bind (bod (Var nam 0)) ((nam, x) : ctx) in
  All nam inp' bod'
bind (Lam nam bod) ctx =
  let bod' = \x -> bind (bod (Var nam 0)) ((nam, x) : ctx) in
  Lam nam bod'
bind (App fun arg) ctx =
  let fun' = bind fun ctx in
  let arg' = bind arg ctx in
  App fun' arg'
bind (Ann chk val typ) ctx =
  let val' = bind val ctx in
  let typ' = bind typ ctx in
  Ann chk val' typ'
bind (Slf nam typ bod) ctx =
  let typ' = bind typ ctx in
  let bod' = \x -> bind (bod (Var nam 0)) ((nam, x) : ctx) in
  Slf nam typ' bod'
bind (Ins val) ctx =
  let val' = bind val ctx in
  Ins val'
bind (ADT scp cts typ) ctx =
  let scp' = map (\x -> bind x ctx) scp in
  let cts' = map (\x -> bindCtr x ctx) cts in
  let typ' = bind typ ctx in
  ADT scp' cts' typ'
  where
    bindCtr (Ctr nm tele)       ctx = Ctr nm (bindTele tele ctx)
    bindTele (TRet term)        ctx = TRet (bind term ctx)
    bindTele (TExt nam typ bod) ctx = TExt nam (bind typ ctx) $ \x -> bindTele (bod x) ((nam, x) : ctx) -- FIXME: 'bod x'?
bind (Con nam arg) ctx =
  let arg' = map (\(f, x) -> (f, bind x ctx)) arg in
  Con nam arg'
bind (Mat cse) ctx =
  let cse' = map (\(cn,cb) -> (cn, bind cb ctx)) cse in
  Mat cse'
bind (Swi zer suc) ctx =
  let zer' = bind zer ctx in
  let suc' = bind suc ctx in
  Swi zer' suc'
bind (Map typ) ctx =
  let typ' = bind typ ctx in
  Map typ'
bind (KVs kvs def) ctx =
  let kvs' = IM.map (\x -> bind x ctx) kvs in
  let def' = bind def ctx in
  KVs kvs' def'
bind (Get g n m k b) ctx =
  let m' = bind m ctx in
  let k' = bind k ctx in
  let b' = \v s -> bind (b v s) ((n, s) : (g, v) : ctx) in
  Get g n m' k' b'
bind (Put g n m k v b) ctx =
  let m' = bind m ctx in
  let k' = bind k ctx in
  let v' = bind v ctx in
  let b' = \o s -> bind (b o s) ((n, s) : (g, o) : ctx) in
  Put g n m' k' v' b'
bind (Ref nam) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> Ref nam
bind (Let nam val bod) ctx =
  let val' = bind val ctx in
  let bod' = \x -> bind (bod (Var nam 0)) ((nam, x) : ctx) in
  Let nam val' bod'
bind (Use nam val bod) ctx =
  let val' = bind val ctx in
  let bod' = \x -> bind (bod (Var nam 0)) ((nam, x) : ctx) in
  Use nam val' bod'
bind Set ctx = Set
bind U64 ctx = U64
bind F64 ctx = F64
bind (Num val) ctx = Num val
bind (Flt val) ctx = Flt val
bind (Op2 opr fst snd) ctx =
  let fst' = bind fst ctx in
  let snd' = bind snd ctx in
  Op2 opr fst' snd'
bind (Txt txt) ctx = Txt txt
bind (Lst lst) ctx =
  let lst' = map (\x -> bind x ctx) lst in
  Lst lst'
bind (Nat val) ctx = Nat val
bind (Hol nam ctxs) ctx = Hol nam (reverse (map snd ctx))
bind (Met uid spn) ctx = Met uid []
bind (Log msg nxt) ctx =
  let msg' = bind msg ctx in
  let nxt' = bind nxt ctx in
  Log msg' nxt'
bind (Var nam idx) ctx =
  case lookup nam ctx of
    Just x  -> x
    Nothing -> Var nam idx
bind (Src src val) ctx =
  let val' = bind val ctx in
  Src src val'

genMetas :: Term -> Term
genMetas term = fst (genMetasGo term 0)

genMetasGo :: Term -> Int -> (Term, Int)
genMetasGo (All nam inp bod) c = 
  let (inp', c1) = genMetasGo inp c
      (bod', c2) = genMetasGo (bod (Var nam 0)) c1
  in (All nam inp' (\_ -> bod'), c2)
genMetasGo (Lam nam bod) c = 
  let (bod', c1) = genMetasGo (bod (Var nam 0)) c
  in (Lam nam (\_ -> bod'), c1)
genMetasGo (App fun arg) c = 
  let (fun', c1) = genMetasGo fun c
      (arg', c2) = genMetasGo arg c1
  in (App fun' arg', c2)
genMetasGo (Ann chk val typ) c = 
  let (val', c1) = genMetasGo val c
      (typ', c2) = genMetasGo typ c1
  in (Ann chk val' typ', c2)
genMetasGo (Slf nam typ bod) c = 
  let (typ', c1) = genMetasGo typ c
      (bod', c2) = genMetasGo (bod (Var nam 0)) c1
  in (Slf nam typ' (\_ -> bod'), c2)
genMetasGo (Ins val) c = 
  let (val', c1) = genMetasGo val c
  in (Ins val', c1)
genMetasGo (ADT scp cts typ) c = 
  let (scp', c1) = foldr (\t (acc, c') -> let (t', c'') = genMetasGo t c' in (t':acc, c'')) ([], c) scp
      (cts', c2) = foldr (\(Ctr nm tele) (acc, c') -> let (tele', c'') = genMetasGoTele tele c' in (Ctr nm tele' : acc, c'')) ([], c1) cts
      (typ', c3) = genMetasGo typ c2
  in (ADT scp' cts' typ', c3)
genMetasGo (Con nam arg) c = 
  let (arg', c1) = foldr (\(f, t) (acc, c') -> let (t', c'') = genMetasGo t c' in ((f, t'):acc, c'')) ([], c) arg
  in (Con nam arg', c1)
genMetasGo (Mat cse) c = 
  let (cse', c1) = foldr (\(cn, cb) (acc, c') -> let (cb', c'') = genMetasGo cb c' in ((cn, cb'):acc, c'')) ([], c) cse
  in (Mat cse', c1)
genMetasGo (Swi zer suc) c = 
  let (zer', c1) = genMetasGo zer c
      (suc', c2) = genMetasGo suc c1
  in (Swi zer' suc', c2)
genMetasGo (Map typ) c = 
  let (typ', c1) = genMetasGo typ c
  in (Map typ', c1)
genMetasGo (KVs kvs def) c = 
  let (def', c1) = genMetasGo def c
      (kvs', c2) = foldr (\ (k, t) (acc, c') -> let (t', c'') = genMetasGo t c' in (IM.insert k t' acc, c'')) (IM.empty, c1) (IM.toList kvs)
  in (KVs kvs' def', c2)
genMetasGo (Get g n m k b) c = 
  let (m', c1) = genMetasGo m c
      (k', c2) = genMetasGo k c1
      (b', c3) = genMetasGo (b (Var g 0) (Var n 0)) c2
  in (Get g n m' k' (\_ _ -> b'), c3)
genMetasGo (Put g n m k v b) c = 
  let (m', c1) = genMetasGo m c
      (k', c2) = genMetasGo k c1
      (v', c3) = genMetasGo v c2
      (b', c4) = genMetasGo (b (Var g 0) (Var n 0)) c3
  in (Put g n m' k' v' (\_ _ -> b'), c4)
genMetasGo (Let nam val bod) c = 
  let (val', c1) = genMetasGo val c
      (bod', c2) = genMetasGo (bod (Var nam 0)) c1
  in (Let nam val' (\_ -> bod'), c2)
genMetasGo (Use nam val bod) c = 
  let (val', c1) = genMetasGo val c
      (bod', c2) = genMetasGo (bod (Var nam 0)) c1
  in (Use nam val' (\_ -> bod'), c2)
genMetasGo (Met _ spn) c = 
  let (spn', c1) = foldr (\t (acc, c') -> let (t', c'') = genMetasGo t c' in (t':acc, c'')) ([], c) spn
  in (Met c1 spn', c1 + 1)
genMetasGo (Op2 opr fst snd) c = 
  let (fst', c1) = genMetasGo fst c
      (snd', c2) = genMetasGo snd c1
  in (Op2 opr fst' snd', c2)
genMetasGo (Lst lst) c = 
  let (lst', c1) = foldr (\t (acc, c') -> let (t', c'') = genMetasGo t c' in (t':acc, c'')) ([], c) lst
  in (Lst lst', c1)
genMetasGo (Log msg nxt) c = 
  let (msg', c1) = genMetasGo msg c
      (nxt', c2) = genMetasGo nxt c1
  in (Log msg' nxt', c2)
genMetasGo (Hol nam ctx) c = 
  let (ctx', c1) = foldr (\t (acc, c') -> let (t', c'') = genMetasGo t c' in (t':acc, c'')) ([], c) ctx
  in (Hol nam ctx', c1)
genMetasGo (Src src val) c = 
  let (val', c1) = genMetasGo val c
  in (Src src val', c1)
genMetasGo term c = (term, c)

genMetasGoTele :: Tele -> Int -> (Tele, Int)
genMetasGoTele (TRet term) c = 
  let (term', c1) = genMetasGo term c
  in (TRet term', c1)
genMetasGoTele (TExt nam typ bod) c = 
  let (typ', c1) = genMetasGo typ c
      (bod', c2) = genMetasGoTele (bod (Var nam 0)) c1
  in (TExt nam typ' (\_ -> bod'), c2)

countMetas :: Term -> Int
countMetas term = snd (genMetasGo term 0)


================================================
FILE: src/Kind/Show.hs
================================================
-- //./Type.hs//

module Kind.Show where

import Prelude hiding (EQ, LT, GT)

import Kind.Type

import Debug.Trace
import Data.Word

import Control.Applicative ((<|>))

import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as IM

-- Stringification
-- ---------------

showTermGo :: Bool -> Term -> Int -> String
showTermGo small term dep =
  case pretty term of
    Just str -> str
    Nothing  -> case term of
      All nam inp bod ->
        let nam' = nam
            inp' = showTermGo small inp dep
            bod' = showTermGo small (bod (Var nam dep)) (dep + 1)
        in concat ["∀(" , nam' , ": " , inp' , ") " , bod']
      Lam nam bod ->
        let nam' = nam
            bod' = showTermGo small (bod (Var nam dep)) (dep + 1)
        in concat ["λ" , nam' , " " , bod']
      App fun arg ->
        let (func, args) = unwrap fun [arg]
            func' = showTermGo small func dep
            args' = unwords (map (\x -> showTermGo small x dep) args)
        in concat ["(" , func' , " " , args' , ")"]
        where unwrap :: Term -> [Term] -> (Term, [Term])
              unwrap (App fun arg) args = unwrap fun (arg:args)
              unwrap term          args = (term, args)
      Ann chk val typ ->
        if small
          then showTermGo small val dep
          else let val' = showTermGo small val dep
                   typ' = showTermGo small typ dep
              in concat ["{" , val' , ": " , typ' , "}"]
      Slf nam typ bod ->
        let nam' = nam
            typ' = showTermGo small typ dep
            bod' = showTermGo small (bod (Var nam dep)) (dep + 1)
        in concat ["$(" , nam' , ": " , typ' , ") " , bod']
      Ins val ->
        let val' = showTermGo small val dep
        in concat ["~" , val']
      -- CHANGED: Updated ADT case to use new Ctr structure
      ADT scp cts typ ->
        let scp' = unwords (map (\x -> showTermGo small x dep) scp)
            cts' = unwords (map (\(Ctr nm tele) -> "#" ++ nm ++ " " ++ showTeleGo small tele dep) cts)
            typ' = showTermGo small typ dep
        in concat ["#[", scp', "]{ ", cts', " } : ", typ']
      Con nam arg ->
        let arg' = unwords (map showArg arg)
        in concat ["#", nam, "{", arg', "}"]
        where
          showArg (maybeField, term) = case maybeField of
            Just field -> field ++ ": " ++ showTermGo small term dep
            Nothing -> showTermGo small term dep
      Mat cse ->
        let cse' = unwords (map (\(cnm, cbod) -> "#" ++ cnm ++ ": " ++ showTermGo small cbod dep) cse)
        in concat ["λ{ ", cse', " }"]
      -- Ref nam -> concat ["@", nam]
      Ref nam -> concat [nam]
      Let nam val bod ->
        let nam' = nam
            val' = showTermGo small val dep
            bod' = showTermGo small (bod (Var nam dep)) (dep + 1)
        in concat ["let " , nam' , " = " , val' , " " , bod']
      Use nam val bod ->
        let nam' = nam
            val' = showTermGo small val dep
            bod' = showTermGo small (bod (Var nam dep)) (dep + 1)
        in concat ["use " , nam' , " = " , val' , " " , bod']
      Set -> "*"
      U64 -> "U64"
      F64 -> "F64"
      Num val ->
        let val' = show val
        in concat [val']
      Flt val ->
        let val' = show val
        in concat [val']
      Op2 opr fst snd ->
        let opr' = showOper opr
            fst' = showTermGo small fst dep
            snd' = showTermGo small snd dep
        in concat ["(" , opr' , " " , fst' , " " , snd' , ")"]
      Swi zero succ ->
        let zero' = showTermGo small zero dep
            succ' = showTermGo small succ dep
        in concat ["λ{ 0: ", zero', " _: ", succ', " }"]
      Map typ ->
        let typ' = showTermGo small typ dep
        in concat ["(Map ", typ', ")"]
      KVs kvs def ->
        let kvs' = unwords (map (\(k, v) -> show k ++ ":" ++ showTermGo small v dep) (IM.toList kvs))
            def' = showTermGo small def dep
        in concat ["{", kvs', " | ", def', "}"]
      Get got nam map key bod ->
        let got' = got
            nam' = nam
            map' = showTermGo small map dep
            key' = showTermGo small key dep
            bod' = showTermGo small (bod (Var got dep) (Var nam dep)) (dep + 2)
        in concat ["get ", got', " = ", nam', "@", map', "[", key', "] ", bod']
      Put got nam map key val bod ->
        let got' = got
            nam' = nam
            map' = showTermGo small map dep
            key' = showTermGo small key dep
            val' = showTermGo small val dep
            bod' = showTermGo small (bod (Var got dep) (Var nam dep)) (dep + 2)
        in concat ["put ", got', " = ", nam', "@", map', "[", key', "] := ", val', " ", bod']
      Txt txt -> concat ["\"" , txt , "\""]
      Lst lst -> concat ["[", unwords (map (\x -> showTermGo small x dep) lst), "]"]
      Nat val -> concat ["#", (show val)]
      Hol nam ctx -> concat ["?" , nam]
      -- Met uid spn -> concat ["_", show uid, "[", strSpn spn dep, " ]"]
      Met uid spn -> concat ["_", show uid]
      Log msg nxt -> 
        let msg' = showTermGo small msg dep
            nxt' = showTermGo small nxt dep
        in concat ["log ", msg', " ", nxt']
      Var nam idx -> nam
      Src src val -> if small
        then showTermGo small val dep
        else concat ["!", showTermGo small val dep]

-- CHANGED: Added showTeleGo function
showTeleGo :: Bool -> Tele -> Int -> String
showTeleGo small tele dep = "{ " ++ go tele dep where
  go (TExt nam typ bod) dep =
    let typ' = showTermGo small typ dep
        bod' = go (bod (Var nam dep)) (dep + 1)
    in concat [nam, ": ", typ', " ", bod']
  go (TRet term) dep =
    let term' = showTermGo small term dep
    in concat ["}: ", term']

showTele :: Tele -> String
showTele tele = showTeleGo True tele 0

showTerm :: Term -> String
showTerm term = showTermGo True term 0

strSpn :: [Term] -> Int -> String
strSpn []       dep = ""
strSpn (x : xs) dep = concat [" ", showTermGo True x dep, strSpn xs dep]

showOper :: Oper -> String
showOper ADD = "+"
showOper SUB = "-"
showOper MUL = "*"
showOper DIV = "/"
showOper MOD = "%"
showOper EQ  = "=="
showOper NE  = "!="
showOper LT  = "<"
showOper GT  = ">"
showOper LTE = "<="
showOper GTE = ">="
showOper AND = "&"
showOper OR  = "|"
showOper XOR = "^"
showOper LSH = "<<"
showOper RSH = ">>"

-- Pretty Printing (Sugars)
-- ------------------------

pretty :: Term -> Maybe String
pretty term = prettyString term <|> prettyNat term <|> prettyList term <|> prettyEqual term

prettyString :: Term -> Maybe String
prettyString (Con "View" [(_, term)]) = do
  chars <- prettyStringGo term
  return $ '"' : chars ++ "\""
prettyString _ = Nothing

prettyStringGo :: Term -> Maybe String
prettyStringGo (Con "Nil" []) = Just []
prettyStringGo (Con "Cons" [(_, Num head), (_, tail)]) = do
  rest <- prettyStringGo tail
  return $ toEnum (fromIntegral head) : rest
prettyStringGo _ = Nothing

prettyNat :: Term -> Maybe String
prettyNat (Con "Zero" []) = Just "#0"
prettyNat term = go 0 term where
  go n (Con "Succ" [(_, pred)]) = go (n + 1) pred
  go n (Con "Zero" []) = Just $ "#" ++ show n
  go _ _ = Nothing

prettyList :: Term -> Maybe String
prettyList term = do
  terms <- asList term
  return $ "[" ++ unwords (map (\x -> showTermGo True x 0) terms) ++ "]"
  where asList (Con "Nil" []) = do
          Just []
        asList (Con "Cons" [(_, head), (_, tail)]) = do
          rest <- asList tail
          return (head : rest)
        asList _ = Nothing

prettyEqual :: Term -> Maybe String
prettyEqual (App (App (App (Ref "Equal") t) a) b) = do
  let a' = showTermGo True a 0
      b' = showTermGo True b 0
  return $ a' ++ " == " ++ b'
prettyEqual _ = Nothing


================================================
FILE: src/Kind/Type.hs
================================================
module Kind.Type where

import System.IO.Unsafe (unsafePerformIO)
import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M

import Debug.Trace
import Data.Word (Word64)

-- Kind's AST
data Term
  -- Product: `∀(x: A) B`
  = All String Term (Term -> Term)

  -- Lambda: `λx f`
  | Lam String (Term -> Term)

  -- Application: `(fun arg)`
  | App Term Term

  -- Annotation: `{x: T}`
  | Ann Bool Term Term

  -- Self-Type: `$(x: A) B`
  | Slf String Term (Term -> Term)

  -- Self-Inst: `~x`
  | Ins Term

  -- Datatype: `#[i0 i1...]{ #C0 Tele0 #C1 Tele1 ... }`
  | ADT [Term] [Ctr] Term

  -- Constructor: `#CN { x0 x1 ... }`
  | Con String [(Maybe String, Term)]

  -- Lambda-Match: `λ{ #C0:B0 #C1:B1 ... }`
  | Mat [(String, Term)]

  -- Top-Level Reference: `Foo`
  | Ref String

  -- Local let-definition: `let x = val body`
  | Let String Term (Term -> Term)

  -- Local use-definition: `use x = val body`
  | Use String Term (Term -> Term)

  -- Universe: `Set`
  | Set

  -- U64 Type: `U64`
  | U64

  -- F64 Type: `F64`
  | F64

  -- U64 Value: `123`
  | Num Word64

  -- F64 Value: `1.5`
  | Flt Double

  -- Binary Operation: `(+ x y)`
  | Op2 Oper Term Term

  -- U64 Elimination: `λ{ 0:A 1+p:B }`
  | Swi Term Term

  -- Linear Map Type: `(Map T)`
  | Map Term 

  -- Linear Map Value: `{ k0:v0 k1:v1 ... | default }`
  | KVs (IM.IntMap Term) Term

  -- Linear Map Getter: `get val = nam@map[key] bod`
  -- - got is the name of the obtained value
  -- - nam is the name of the map
  -- - map is the value of the map
  -- - key is the key to query
  -- - bod is the continuation; receives the value and the same map
  | Get String String Term Term (Term -> Term -> Term)

  -- Map Swapper: `put got = nam@map[key] := val body`
  -- - got is the name of the old value
  -- - nam is the name of the map
  -- - map is the value of the map
  -- - key is the key to swap
  -- - val is the val to insert
  -- - bod is the continuation; receives the old value and the changed map
  | Put String String Term Term Term (Term -> Term -> Term)

  -- Inspection Hole
  | Hol String [Term]

  -- Unification Metavar
  | Met Int [Term]

  -- Logging
  | Log Term Term

  -- Variable
  | Var String Int

  -- Source Location
  | Src Cod Term

  -- Text Literal (sugar)
  | Txt String

  -- List Literal (sugar)
  | Lst [Term]

  -- Nat Literal (sugar)
  | Nat Integer

  -- Substitution
  | Sub Term

-- Location: Name, Line, Column
data Loc = Loc String Int Int
data Cod = Cod Loc Loc

-- Numeric Operators
data Oper 
  = ADD | SUB | MUL | DIV
  | MOD | EQ  | NE  | LT
  | GT  | LTE | GTE | AND
  | OR  | XOR | LSH | RSH
  deriving Show

-- Telescope
data Tele
  = TRet Term
  | TExt String Term (Term -> Tele)

-- Constructor
data Ctr = Ctr String Tele

-- Book of Definitions
type Book = M.Map String Term

-- Type-Checker Outputs
data Info
  = Found String Term [Term] Int
  | Solve Int Term Int
  | Error (Maybe Cod) Term Term Term Int
  | Vague String
  | Print Term Int

-- Unification Solutions
type Fill = IM.IntMap Term

-- Checker State
data Check = Check (Maybe Cod) Term Term Int -- postponed check
data State = State Book Fill [Check] [Info] -- state type
data Res a = Done State a | Fail State -- result type
data Env a = Env (State -> Res a) -- monadic checker

-- UNCOMMENT THIS TO DEBUG THE TYPE CHECKER
-- debug a b = trace a b
debug a b = b


================================================
FILE: src/Kind/Util.hs
================================================
-- //./Type.hs//

module Kind.Util where

import Kind.Show
import Kind.Type
import Kind.Equal

import Prelude hiding (LT, GT, EQ)

import qualified Data.IntMap.Strict as IM
import qualified Data.Map.Strict as M
import qualified Data.Set as S

import Debug.Trace

-- Gets dependencies of a term
getDeps :: Term -> [String]
getDeps term = case term of
  Ref nam         -> [nam]
  All _ inp out   -> getDeps inp ++ getDeps (out Set)
  Lam _ bod       -> getDeps (bod Set)
  App fun arg     -> getDeps fun ++ getDeps arg
  Ann _ val typ   -> getDeps val ++ getDeps typ
  Slf _ typ bod   -> getDeps typ ++ getDeps (bod Set)
  Ins val         -> getDeps val
  ADT scp cts t   -> concatMap getDeps scp ++ concatMap getDepsCtr cts ++ getDeps t
  Con _ arg       -> concatMap (getDeps . snd) arg
  Mat cse         -> concatMap (getDeps . snd) cse
  Let _ val bod   -> getDeps val ++ getDeps (bod Set)
  Use _ val bod   -> getDeps val ++ getDeps (bod Set)
  Op2 _ fst snd   -> getDeps fst ++ getDeps snd
  Swi zer suc     -> getDeps zer ++ getDeps suc
  Map val         -> getDeps val
  KVs kvs def     -> concatMap getDeps (IM.elems kvs) ++ getDeps def
  Get _ _ m k b   -> getDeps m ++ getDeps k ++ getDeps (b Set Set)
  Put _ _ m k v b -> getDeps m ++ getDeps k ++ getDeps v ++ getDeps (b Set Set)
  Src _ val       -> getDeps val
  Hol _ args      -> concatMap getDeps args
  Met _ args      -> concatMap getDeps args
  Log msg nxt     -> getDeps msg ++ getDeps nxt
  Var _ _         -> []
  Set             -> []
  U64             -> []
  F64             -> []
  Num _           -> []
  Flt _           -> []
  Txt _           -> []
  Lst elems       -> concatMap getDeps elems
  Nat _           -> []

-- Gets dependencies of a constructor
getDepsCtr :: Ctr -> [String]
getDepsCtr (Ctr _ tele) = getDepsTele tele

-- Gets dependencies of a telescope
getDepsTele :: Tele -> [String]
getDepsTele (TRet term) = getDeps term
getDepsTele (TExt _ typ bod) = getDeps typ ++ getDepsTele (bod Set)

-- Gets all dependencies (direct and indirect) of a term
getAllDeps :: Book -> String -> S.Set String
getAllDeps book name = go S.empty [name] where
  go visited [] = visited
  go visited (x:xs)
    | S.member x visited = go visited xs
    | otherwise = case M.lookup x book of
        Just term -> go (S.insert x visited) (getDeps term ++ xs)
        Nothing   -> go (S.insert x visited) xs

-- Topologically sorts a book
topoSortBook :: Book -> [(String, Term)]
topoSortBook book = go (M.keysSet book) [] where
  go mustInclude done = case S.lookupMin mustInclude of
    Nothing   -> reverse done
    Just name -> 
      let (mustInclude', done') = include mustInclude done name
      in go mustInclude' done'

  include :: S.Set String -> [(String, Term)] -> String -> (S.Set String, [(String, Term)])
  include mustInclude done name =
    if not (S.member name mustInclude) then
      (mustInclude, done)
    else case M.lookup name book of
      Nothing ->
        error ("unbound:" ++ name)
      Just term ->
        let deps = getDeps term
            mustInclude' = S.delete name mustInclude
            (mustInclude'', done') = includeDeps mustInclude' done deps
        in (mustInclude'', (name,term) : done')

  includeDeps :: S.Set String -> [(String, Term)] -> [String] -> (S.Set String, [(String, Term)])
  includeDeps mustInclude done [] = (mustInclude, done)
  includeDeps mustInclude done (dep:deps) =
    let (mustInclude', done') = include mustInclude done dep
        (mustInclude'', done'') = includeDeps mustInclude' done' deps
    in (mustInclude'', done'')

-- Converts:
-- - from a Tele: `{ x:A y:(B x) ... }: (C x y ...)`
-- - to a type: `∀(x: A) ∀(y: (B x)) ... (C x y ...)`
teleToType :: Tele -> Term -> Int -> Term
teleToType (TRet _)           ret _   = ret
teleToType (TExt nam inp bod) ret dep = All nam inp (\x -> teleToType (bod x) ret (dep + 1))

-- Converts:
-- - from a Tele : `{ x:A y:(B x) ... }: (C x y ...)`
-- - to terms    : `([(Just "x", <A>), [(Just "y", <(B x)>)], ...], <(C x y ...)>)`
teleToTerms :: Tele -> Int -> ([(Maybe String, Term)], Term)
teleToTerms tele dep = go tele [] dep where
  go (TRet ret)         args _   = (reverse args, ret)
  go (TExt nam inp bod) args dep = go (bod (Var nam dep)) ((Just nam, Var nam dep) : args) (dep + 1)

getTeleNames :: Tele -> Int -> [String] -> [String]
getTeleNames (TRet _)           dep acc = reverse acc
getTeleNames (TExt name _ next) dep acc = getTeleNames (next (Var name dep)) (dep+1) (name:acc)

getTeleFields :: Tele -> Int -> [(String,Term)] -> [(String,Term)]
getTeleFields (TRet _)              dep acc = reverse acc
getTeleFields (TExt name ttyp next) dep acc = getTeleFields (next (Var name dep)) (dep+1) ((name,ttyp):acc)

getDatIndices :: Term -> [Term]
getDatIndices term = case term of
  ADT idxs _ _ -> idxs
  _            -> []

getType :: Term -> Term
getType (Ann _ val typ) = typ
getType _               = error "?"

getTerm :: Term -> Term
getTerm (Ann _ val typ) = val
getTerm _               = error "?"

getCtrName :: Ctr -> String
getCtrName (Ctr name _) = name

getADTCts :: Term -> [(String,Ctr)]
getADTCts (ADT _ cts _) = map (\ ctr -> (getCtrName ctr, ctr)) cts
getADTCts (Src loc val) = getADTCts val
getADTCts term          = error ("not-an-adt:" ++ showTerm term)

-- Given a typed term, return its argument's names
getArgNames :: Term -> [String]
getArgNames (Ann _ _ typ) = getForallNames typ
getArgNames (Src _ val)   = getArgNames val
getArgNames _             = []

-- Returns the names in a chain of foralls
getForallNames :: Term -> [String]
getForallNames (All nam _ bod) = nam : getForallNames (bod Set)
getForallNames (Src _ val)     = getForallNames val
getForallNames _               = []

getOpReturnType :: Oper -> Term -> Term
getOpReturnType ADD U64 = U64
getOpReturnType ADD F64 = F64
getOpReturnType SUB U64 = U64
getOpReturnType SUB F64 = F64
getOpReturnType MUL U64 = U64
getOpReturnType MUL F64 = F64
getOpReturnType DIV U64 = U64
getOpReturnType DIV F64 = F64
getOpReturnType MOD U64 = U64
getOpReturnType EQ  _   = U64
getOpReturnType NE  _   = U64
getOpReturnType LT  _   = U64
getOpReturnType GT  _   = U64
getOpReturnType LTE _   = U64
getOpReturnType GTE _   = U64
getOpReturnType AND U64 = U64
getOpReturnType OR  U64 = U64
getOpReturnType XOR U64 = U64
getOpReturnType LSH U64 = U64
getOpReturnType RSH U64 = U64
getOpReturnType opr trm = error ("Invalid opertor: " ++ (show opr) ++ " Invalid operand type: " ++ (showTerm trm))

checkValidType :: Term -> [Term] -> Int -> Env Bool
checkValidType typ validTypes dep = foldr (\t acc -> do
    isEqual <- equal typ t dep
    if isEqual then return True else acc
  ) (return False) validTypes



================================================
FILE: src/Kind.hs
================================================
module Kind (
  module Kind.CLI,
  module Kind.Check,
  module Kind.Env,
  module Kind.Equal,
  module Kind.Parse,
  module Kind.Reduce,
  module Kind.Show,
  module Kind.Type,
  module Kind.Util,
) where

import Kind.CLI
import Kind.Check
import Kind.CompileJS
import Kind.Env
import Kind.Equal
import Kind.Parse
import Kind.Reduce
import Kind.Show
import Kind.Type
import Kind.Util