summaryrefslogtreecommitdiff
path: root/src/LambdaCube/Compiler/DeBruijn.hs
blob: f929f6b7041ccc28935275101aff2e141bc431e3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module LambdaCube.Compiler.DeBruijn where

import Data.Binary
import GHC.Generics (Generic)

import Data.Bits
import Control.Arrow hiding ((<+>))

import LambdaCube.Compiler.Utils
import LambdaCube.Compiler.Pretty

-------------------------------------------------------------------------------- rearrange De Bruijn indices

class Rearrange a where
    rearrange :: Int{-level-} -> RearrangeFun -> a -> a

data RearrangeFun
    = RFSubst Int Int
    | RFMove Int
    | RFUp Int
    deriving Show

rearrangeFun = \case
    RFSubst i j -> \k -> if k == i then j else if k > i then k - 1 else k
    RFMove i -> \k -> if k == i then 0 else k + 1
    RFUp n -> \k -> if k >= 0 then k + n else k

rSubst :: Rearrange a => Int -> Int -> a -> a
rSubst i j = rearrange 0 $ RFSubst i j

-- move index to 0
rMove :: Rearrange a => Int -> Int -> a -> a
rMove i l = rearrange l $ RFMove i

rUp :: Rearrange a => Int -> Int -> a -> a
rUp n l = rearrange l $ RFUp n

up1_ :: Rearrange a => Int -> a -> a
up1_ = rUp 1

up n = rUp n 0
up1 = up1_ 0

instance Rearrange a => Rearrange [a] where
    rearrange l f = map $ rearrange l f

instance (Rearrange a, Rearrange b) => Rearrange (Either a b) where
    rearrange l f = rearrange l f +++ rearrange l f

instance (Rearrange a, Rearrange b) => Rearrange (a, b) where
    rearrange l f = rearrange l f *** rearrange l f

instance Rearrange Void where
    rearrange _ _ = elimVoid

------------------------------------------------------- set of free variables (implemented with bit vectors)

newtype FreeVars = FreeVars Integer
    deriving (Eq, Generic, Show)

instance Binary FreeVars

instance PShow FreeVars where
    pShow (FreeVars s) = "FreeVars" `DApp` pShow s

instance Monoid FreeVars where
    mempty = FreeVars 0
#if !MIN_VERSION_base(4,11,0)
    FreeVars a `mappend` FreeVars b = FreeVars $ a .|. b
#else
instance Semigroup FreeVars where
    FreeVars a <> FreeVars b = FreeVars $ a .|. b
#endif

freeVar :: Int -> FreeVars
freeVar i = FreeVars $ 1 `shiftL` i

delVar :: Int -> FreeVars -> FreeVars
delVar 0 (FreeVars i) = FreeVars $ i `shiftR` 1
delVar 1 (FreeVars i) = FreeVars $ if testBit i 0 then (i `shiftR` 1) `setBit` 0 else (i `shiftR` 1) `clearBit` 0 
delVar l (FreeVars i) = FreeVars $ case i `shiftR` (l+1) of
    0 -> i `clearBit` l
    x -> (x `shiftL` l) .|. (i .&. ((1 `shiftL` l)-1))

shiftFreeVars :: Int -> FreeVars -> FreeVars
shiftFreeVars i (FreeVars x) = FreeVars $ x `shift` i

usedVar :: HasFreeVars a => Int -> a -> Bool
usedVar i (getFreeVars -> FreeVars a) = testBit a i

freeVars :: FreeVars -> [Int]
freeVars (FreeVars x) = take (popCount x) [i | i <- [0..], testBit x i]

isClosed :: FreeVars -> Bool
isClosed (FreeVars x) = x == 0

lowerFreeVars = shiftFreeVars (-1)

rearrangeFreeVars g l (FreeVars i) = FreeVars $ case g of
    RFUp n -> ((i `shiftR` l) `shiftL` (n+l)) .|. (i .&. ((1 `shiftL` l)-1))
    RFMove n -> (f $ (i `shiftR` l) `shiftL` (l+1)) .|. (i .&. ((1 `shiftL` l)-1))
      where
        f x = if testBit x (n+l+1) then x `clearBit` (n+l+1) `setBit` l else x
    _ -> error $ "rearrangeFreeVars: " ++ show g


-- TODO: rename
dbGE i x = dbGE_ i $ getFreeVars x

dbGE_ i (FreeVars x) = (1 `shiftL` i) > x

-------------------------------------------------------------------------------- type class for getting free variables

class HasFreeVars a where
    getFreeVars :: a -> FreeVars

instance HasFreeVars FreeVars where
    getFreeVars = id

instance HasFreeVars Void where
    getFreeVars = elimVoid

-------------------------------------------------------------------------------- substitute names with De Bruijn indices

class DeBruijnify n a where
    deBruijnify_ :: Int -> [n] -> a -> a

deBruijnify = deBruijnify_ 0

instance (DeBruijnify n a, DeBruijnify n b) => DeBruijnify n (a, b) where
    deBruijnify_ k ns = deBruijnify_ k ns *** deBruijnify_ k ns

instance (DeBruijnify n a, DeBruijnify n b) => DeBruijnify n (Either a b) where
    deBruijnify_ k ns = deBruijnify_ k ns +++ deBruijnify_ k ns

instance (DeBruijnify n a) => DeBruijnify n [a] where
    deBruijnify_ k ns = fmap (deBruijnify_ k ns)