
   0 | module Compiler.Separate
   1 | 
   2 | import public Core.FC
   3 | import public Core.Name
   4 | import public Core.Name.Namespace
   5 | import public Core.CompileExpr
   6 | import public Compiler.LambdaLift
   7 | import public Compiler.VMCode
   8 | import public Libraries.Data.Graph
   9 | import public Data.SortedMap
  10 | import public Data.SortedSet
  11 | import public Libraries.Data.StringMap
  12 | 
  13 | import Core.Hash
  14 | import Core.TT
  15 | import Data.List
  16 | import Data.List1
  17 | import Data.Vect
  18 | import Data.Maybe
  19 | 
  20 | %default covering
  21 | 
  22 | -- Compilation unit IDs are intended to be opaque,
  23 | -- just to be able to express dependencies via keys in a map and such.
  24 | export
  25 | record CompilationUnitId where
  26 |   constructor CUID
  27 |   int : Int
  28 | 
  29 | export
  30 | Eq CompilationUnitId where
  31 |   CUID x == CUID y = x == y
  32 | 
  33 | export
  34 | Ord CompilationUnitId where
  35 |   compare (CUID x) (CUID y) = compare x y
  36 | 
  37 | export
  38 | Hashable CompilationUnitId where
  39 |   hashWithSalt h (CUID int) = hashWithSalt h int
  40 | 
  41 | ||| A compilation unit is a set of namespaces.
  42 | |||
  43 | ||| The record is parameterised by the type of the definition,
  44 | ||| which makes it reusable for various IRs provided by getCompileData.
  45 | public export
  46 | record CompilationUnit def where
  47 |   constructor MkCompilationUnit
  48 | 
  49 |   ||| Unique identifier of a compilation unit within a CompilationUnitInfo record.
  50 |   id : CompilationUnitId
  51 | 
  52 |   ||| Namespaces contained within the compilation unit.
  53 |   namespaces : List1 Namespace
  54 | 
  55 |   ||| Other units that this unit depends on.
  56 |   dependencies : SortedSet CompilationUnitId
  57 | 
  58 |   ||| The definitions belonging into this compilation unit.
  59 |   definitions : List (Name, def)
  60 | 
  61 | export
  62 | Hashable def => Hashable (CompilationUnit def) where
  63 |   hashWithSalt h cu =
  64 |     h `hashWithSalt` Prelude.toList cu.dependencies
  65 |       `hashWithSalt` cu.definitions
  66 | 
  67 | private
  68 | getNS : Name -> Namespace
  69 | getNS (NS ns _) = ns
  70 | getNS _ = emptyNS
  71 | 
  72 | ||| Group definitions by namespace.
  73 | private
  74 | splitByNS : List (Name, def) -> List (Namespace, List (Name, def))
  75 | splitByNS = SortedMap.toList . foldl addOne SortedMap.empty
  76 |   where
  77 |     addOne
  78 |       : SortedMap Namespace (List (Name, def))
  79 |       -> (Name, def)
  80 |       -> SortedMap Namespace (List (Name, def))
  81 |     addOne nss ndef@(n, _) =
  82 |       SortedMap.mergeWith
  83 |         (++)
  84 |         (SortedMap.singleton (getNS n) [ndef])
  85 |         nss
  86 | 
  87 | public export
  88 | interface HasNamespaces a where
  89 |   ||| Return the set of namespaces mentioned within
  90 |   nsRefs : a -> SortedSet Namespace
  91 | 
  92 | -- For now, we have instances only for NamedDef, LiftedDef and VMDef.
  93 | -- For other IR representations, we'll have to add more instances.
  94 | -- This is not hard, just a bit of tedious mechanical work.
  95 | mutual
  96 |   export
  97 |   HasNamespaces NamedCExp where
  98 |     nsRefs (NmLocal fc n) = SortedSet.empty
  99 |     nsRefs (NmRef fc n) = SortedSet.singleton $ getNS n
 100 |     nsRefs (NmLam fc n rhs) = nsRefs rhs
 101 |     nsRefs (NmLet fc n val rhs) = nsRefs val <+> nsRefs rhs
 102 |     nsRefs (NmApp fc f args) = nsRefs f <+> concatMap nsRefs args
 103 |     nsRefs (NmCon fc cn ci tag args) = concatMap nsRefs args
 104 |     nsRefs (NmForce fc reason rhs) = nsRefs rhs
 105 |     nsRefs (NmDelay fc reason rhs) = nsRefs rhs
 106 |     nsRefs (NmErased fc) = SortedSet.empty
 107 |     nsRefs (NmPrimVal fc x) = SortedSet.empty
 108 |     nsRefs (NmOp fc op args) = concatMap nsRefs args
 109 |     nsRefs (NmExtPrim fc n args) = concatMap nsRefs args
 110 |     nsRefs (NmConCase fc scrut alts mbDflt) =
 111 |       nsRefs scrut <+> concatMap nsRefs alts <+> concatMap nsRefs mbDflt
 112 |     nsRefs (NmConstCase fc scrut alts mbDflt) =
 113 |       nsRefs scrut <+> concatMap nsRefs alts <+> concatMap nsRefs mbDflt
 114 |     nsRefs (NmCrash fc msg) = SortedSet.empty
 115 | 
 116 |   export
 117 |   HasNamespaces NamedConAlt where
 118 |     nsRefs (MkNConAlt n ci tag args rhs) = nsRefs rhs
 119 | 
 120 |   export
 121 |   HasNamespaces NamedConstAlt where
 122 |     nsRefs (MkNConstAlt c rhs) = nsRefs rhs
 123 | 
 124 |   export
 125 |   HasNamespaces NamedDef where
 126 |     nsRefs (MkNmFun argNs rhs) = nsRefs rhs
 127 |     nsRefs (MkNmCon tag arity nt) = SortedSet.empty
 128 |     nsRefs (MkNmForeign ccs fargs rty) = SortedSet.empty
 129 |     nsRefs (MkNmError rhs) = nsRefs rhs
 130 | 
 131 | mutual
 132 |   export
 133 |   HasNamespaces (Lifted vars) where
 134 |     nsRefs (LLocal fc prf) = SortedSet.empty
 135 |     nsRefs (LAppName fc reason n args) =
 136 |       SortedSet.singleton (getNS n) <+> concatMap nsRefs args
 137 |     nsRefs (LUnderApp fc n missing args) =
 138 |       SortedSet.singleton (getNS n) <+> concatMap nsRefs args
 139 |     nsRefs (LApp fc reason f args) = nsRefs f <+> nsRefs args
 140 |     nsRefs (LLet fc n val rhs) = nsRefs val <+> nsRefs rhs
 141 |     nsRefs (LCon fc cn ci tag args) = concatMap nsRefs args
 142 |     nsRefs (LOp fc reason op args) = concatMap nsRefs args
 143 |     nsRefs (LExtPrim fc reason n args) = concatMap nsRefs args
 144 |     nsRefs (LConCase fc scrut alts mbDflt) =
 145 |       nsRefs scrut <+> concatMap nsRefs alts <+> concatMap nsRefs mbDflt
 146 |     nsRefs (LConstCase fc scrut alts mbDflt) =
 147 |       nsRefs scrut <+> concatMap nsRefs alts <+> concatMap nsRefs mbDflt
 148 |     nsRefs (LPrimVal fc x) = SortedSet.empty
 149 |     nsRefs (LErased fc) = SortedSet.empty
 150 |     nsRefs (LCrash fc msg) = SortedSet.empty
 151 | 
 152 |   export
 153 |   HasNamespaces (LiftedConAlt vars) where
 154 |     nsRefs (MkLConAlt n ci tag args rhs) = nsRefs rhs
 155 | 
 156 |   export
 157 |   HasNamespaces (LiftedConstAlt vars) where
 158 |     nsRefs (MkLConstAlt c rhs) = nsRefs rhs
 159 | 
 160 |   export
 161 |   HasNamespaces LiftedDef where
 162 |     nsRefs (MkLFun args scope rhs) = nsRefs rhs
 163 |     nsRefs (MkLCon tag arity nt) = SortedSet.empty
 164 |     nsRefs (MkLForeign ccs fargs rty) = SortedSet.empty
 165 |     nsRefs (MkLError rhs) = nsRefs rhs
 166 | 
 167 | export
 168 | HasNamespaces VMInst where
 169 |   nsRefs (DECLARE x) = empty
 170 |   nsRefs START = empty
 171 |   nsRefs (ASSIGN x y) = empty
 172 |   nsRefs (MKCON x tag args) = either (const empty) (singleton . getNS) tag
 173 |   nsRefs (MKCLOSURE x n missing args) = singleton $ getNS n
 174 |   nsRefs (MKCONSTANT x y) = empty
 175 |   nsRefs (APPLY x f a) = empty
 176 |   nsRefs (CALL x tailpos n args) = singleton $ getNS n
 177 |   nsRefs (OP x y xs) = empty
 178 |   nsRefs (EXTPRIM x n xs) = singleton $ getNS n
 179 |   nsRefs (CASE x alts def) =
 180 |     maybe empty (concatMap nsRefs) def <+>
 181 |     concatMap ((concatMap nsRefs) . snd) alts <+>
 182 |     concatMap ((either (const empty) (singleton . getNS)) . fst) alts
 183 |   nsRefs (CONSTCASE x alts def) =
 184 |     maybe empty (concatMap nsRefs) def <+>
 185 |     concatMap ((concatMap nsRefs) . snd) alts
 186 |   nsRefs (PROJECT x value pos) = empty
 187 |   nsRefs (NULL x) = empty
 188 |   nsRefs (ERROR x) = empty
 189 | 
 190 | export
 191 | HasNamespaces VMDef where
 192 |   nsRefs (MkVMFun args is) = concatMap nsRefs is
 193 |   nsRefs (MkVMForeign {}) = empty
 194 |   nsRefs (MkVMError is) = concatMap nsRefs is
 195 | 
 196 | 
 197 | -- a slight hack for convenient use with CompileData.namedDefs
 198 | export
 199 | HasNamespaces a => HasNamespaces (FC, a) where
 200 |   nsRefs (_, x) = nsRefs x
 201 | 
 202 | -- another slight hack for convenient use with CompileData.namedDefs
 203 | export
 204 | Hashable def => Hashable (FC, def) where
 205 |   -- ignore FC in hash, like everywhere else
 206 |   hashWithSalt h (fc, x) = hashWithSalt h x
 207 | 
 208 | ||| Output of the codegen separation algorithm.
 209 | ||| Should contain everything you need in a separately compiling codegen.
 210 | public export
 211 | record CompilationUnitInfo def where
 212 |   constructor MkCompilationUnitInfo
 213 | 
 214 |   ||| Compilation units computed from the given definitions,
 215 |   ||| ordered topologically, starting from units depending on no other unit.
 216 |   compilationUnits : List (CompilationUnit def)
 217 | 
 218 |   ||| Mapping from ID to CompilationUnit.
 219 |   byId : SortedMap CompilationUnitId (CompilationUnit def)
 220 | 
 221 |   ||| Maps each namespace to the compilation unit that contains it.
 222 |   namespaceMap : SortedMap Namespace CompilationUnitId
 223 | 
 224 | ||| Group the given definitions into compilation units for separate code generation.
 225 | export
 226 | getCompilationUnits : HasNamespaces def => List (Name, def) -> CompilationUnitInfo def
 227 | getCompilationUnits {def} defs =
 228 |   let
 229 |     -- Definitions grouped by namespace.
 230 |     defsByNS : SortedMap Namespace (List (Name, def))
 231 |       = SortedMap.fromList $ splitByNS defs
 232 | 
 233 |     -- Mapping from a namespace to all namespaces mentioned within.
 234 |     -- Represents graph edges pointing in that direction.
 235 |     nsDeps : SortedMap Namespace (SortedSet Namespace)
 236 |       = foldl (SortedMap.mergeWith SortedSet.union) SortedMap.empty
 237 |           [ SortedMap.singleton (getNS n) (SortedSet.delete (getNS n) (nsRefs d))
 238 |           | (n, d) <- defs
 239 |           ]
 240 | 
 241 |     -- Strongly connected components of the NS dep graph,
 242 |     -- ordered by output degree ascending.
 243 |     --
 244 |     -- Each SCC will become a compilation unit.
 245 |     components : List (List1 Namespace)
 246 |       = List.reverse $ tarjan nsDeps  -- tarjan generates reverse toposort
 247 | 
 248 |     -- Maps a namespace to the compilation unit that contains it.
 249 |     nsMap : SortedMap Namespace CompilationUnitId
 250 |       = SortedMap.fromList [(ns, cuid) | (cuid, nss) <- withCUID components, ns <- List1.forget nss]
 251 | 
 252 |     -- List of all compilation units, ordered by number of dependencies, ascending.
 253 |     units : List (CompilationUnit def)
 254 |       = [mkUnit nsDeps nsMap defsByNS cuid nss | (cuid, nss) <- withCUID components]
 255 | 
 256 |   in MkCompilationUnitInfo
 257 |       { compilationUnits = units
 258 |       , byId = SortedMap.fromList [(unit.id, unit) | unit <- units]
 259 |       , namespaceMap = nsMap
 260 |       }
 261 | 
 262 |   where
 263 |     withCUID : List a -> List (CompilationUnitId, a)
 264 |     withCUID xs = [(CUID $ cast i, x) | (i, x) <- zip [0..length xs] xs]
 265 | 
 266 |     ||| Wrap all information in a compilation unit record.
 267 |     mkUnit :
 268 |       SortedMap Namespace (SortedSet Namespace)
 269 |       -> SortedMap Namespace CompilationUnitId
 270 |       -> SortedMap Namespace (List (Name, def))
 271 |       -> CompilationUnitId -> List1 Namespace -> CompilationUnit def
 272 |     mkUnit nsDeps nsMap defsByNS cuid nss =
 273 |       MkCompilationUnit
 274 |       { id = cuid
 275 |       , namespaces = nss
 276 |       , dependencies = SortedSet.delete cuid dependencies
 277 |       , definitions = definitions
 278 |       }
 279 |      where
 280 |       dependencies : SortedSet CompilationUnitId
 281 |       dependencies = SortedSet.fromList $ do
 282 |         ns <- List1.forget nss  -- NS contained within
 283 |         depsNS <- Prelude.toList $  -- NS we depend on
 284 |           fromMaybe SortedSet.empty $
 285 |             SortedMap.lookup ns nsDeps
 286 | 
 287 |         case SortedMap.lookup depsNS nsMap of
 288 |           Nothing => []
 289 |           Just depCUID => [depCUID]
 290 | 
 291 |       definitions : List (Name, def)
 292 |       definitions = concat [fromMaybe [] $ SortedMap.lookup ns defsByNS | ns <- nss]
 293 |