0 | module Syntax.StringDiagram.Arrow

  2 | import public Control.Monad.State

  3 | import public Control.Applicative.Const

  4 | import public Control.Arrow

  5 | import public Control.Category.Core

  6 | import public Language.Reflection

  7 | import public Data.Wrap0

  9 | %default total

 10 | %language ElabReflection

 12 | export infix 0 -<

 13 | export prefix 0 =<

 16 | record ArrowVar where

 17 |   constructor MkAVar

 18 |   name : String

 19 |   type : TTImp

 22 | zipFilter : List a -> List Bool -> List a

 23 | zipFilter = catMaybes .: zipWith (\x,b => guard b $> x)

 26 | bindVars : TTImp -> List String

 27 | bindVars = (nub . runConst .) $ mapATTImp' $ \case

 28 |   IBindVar _ (UN $ Basic var) => const $ MkConst [var]

 29 |   IAs _ _ _ (UN $ Basic var) _ => \(MkConst vs) => MkConst (var :: vs)

 30 |   _ => id

 32 | isBoundVar : TTImp -> String -> Bool

 33 | isBoundVar t var =

 34 |   let _ = Lazy.Monoid.Any

 35 |   in runConst {a=Lazy Bool} $ flip mapATTImp' t $ \case

 36 |     IBindVar _ (UN $ Basic var') => const $ MkConst $ delay (var == var')

 37 |     IAs _ _ _ (UN $ Basic var') _ => \(MkConst b) => MkConst (var == var' || b)

 38 |     _ => id

 40 | usedInExpr : TTImp -> String -> Bool

 41 | usedInExpr t var =

 42 |   let _ = Lazy.Monoid.Any

 43 |   in runConst {a=Lazy Bool} $ flip mapATTImp' t $ \case

 44 |     -- If a hole is found, always include every possible variable

 45 |     IHole _ _ => const $ MkConst $ delay True

 46 |     IVar _ (UN $ Basic var') => const $ MkConst $ delay (var == var')

 47 |     _ => id

 49 | usedInDiagram : TTImp -> String -> Bool

 50 | usedInDiagram (ILam _ _ _ (Just $ UN $ Basic v) _ rest) var =

 51 |   v /= var && usedInDiagram rest var

 52 | usedInDiagram (ILam _ _ _ _ _ rest) var = usedInDiagram rest var

 53 | usedInDiagram (ILet _ _ _ (UN $ Basic v) _ exp rest) var =

 54 |   usedInExpr exp var || (v /= var && usedInDiagram rest var)

 55 | usedInDiagram (ILet _ _ _ _ _ exp rest) var =

 56 |   usedInExpr exp var || usedInDiagram rest var

 57 | usedInDiagram (ICase _ _ exp _ clauses) var =

 58 |   usedInExpr exp var || any (\case

 59 |     PatClause _ lhs rhs =>

 60 |       not (isBoundVar lhs var) &&

 61 |         usedInDiagram (assert_smaller clauses rhs) var

 62 |     _ => False

 63 |     ) clauses

 64 | usedInDiagram `((~(mor) -< ~(inp)) >>= ~(rest)) var =

 65 |   usedInExpr inp var || usedInDiagram rest var

 66 | usedInDiagram `((~(mor) -< ~(inp)) >> ~(rest)) var =

 67 |   usedInExpr inp var || usedInDiagram rest var

 68 | usedInDiagram `(=< ~(out)) var = usedInExpr out var

 69 | usedInDiagram _ _ = False

 73 | arrowImpl : TTImp -> Elab TTImp

 74 | arrowImpl t = do

 75 |   ts <- arrowImplLam [<] [] t

 76 |   pure $ composeImp $ optimize ts

 77 |   where

 78 |     optimize : SnocList TTImp -> SnocList TTImp

 79 |     optimize [<] = [<]

 80 |     optimize [<t] = [<t]

 81 |     optimize (ts :< `(Control.Arrow.arrow {a = ~a, b = ~_} ~f) :< `(Control.Arrow.arrow {a = ~b, b = ~c} ~g)) =

 82 |       assert_total $ optimize (ts :< `(Control.Arrow.arrow {a = ~a, b = ~c} (Prelude.(.) {a = ~a, b = ~b, c = ~c} ~g ~f)))

 83 |     optimize (ts :< t :< t') = optimize (ts :< t) :< t'

 85 |     composeImp : SnocList TTImp -> TTImp

 86 |     composeImp [<] = `(Control.Category.Core.id)

 87 |     composeImp [<t] = t

 88 |     composeImp (ts :< t) =

 89 |       `(Control.Category.Core.(.)

 90 |         ~t ~(composeImp ts))

 92 |     pairTy : List ArrowVar -> TTImp

 93 |     pairTy [] = `(Builtin.Unit)

 94 |     pairTy [MkAVar _ ty] = ty

 95 |     pairTy (MkAVar _ ty :: vars) = `(Builtin.Pair ~ty ~(pairTy vars))

 97 |     pair : List String -> TTImp

 98 |     pair [] = `(Builtin.MkUnit)

 99 |     pair [n] = IVar EmptyFC $ UN $ Basic n

100 |     pair (n :: ns) = `(Builtin.MkPair ~(IVar EmptyFC $ UN $ Basic n) ~(pair ns))

102 |     stringBind : List (ArrowVar, Bool) -> TTImp

103 |     stringBind [] = `(Builtin.MkUnit)

104 |     stringBind [(v,b)] = if b then IBindVar EmptyFC (UN $ Basic v.name) else Implicit EmptyFC True

105 |     stringBind ((v,b) :: strings) = `(Builtin.MkPair ~(if b then IBindVar EmptyFC (UN $ Basic v.name) else Implicit EmptyFC True) ~(stringBind strings))

107 |     stringLam : List (ArrowVar, Bool) -> TTImp -> TTImp

108 |     stringLam [] t = `(\_ => ~t)

109 |     stringLam [(_,False)] t = `(\_ => ~t)

110 |     stringLam [(v,True)] t = ILam EmptyFC MW ExplicitArg (Just $ UN $ Basic v.name) (Implicit EmptyFC False) t

111 |     stringLam ns t = `(\ ~(stringBind ns) => ~t)

113 |     stringLamMN : Name -> List (ArrowVar, Bool) -> TTImp -> TTImp

114 |     stringLamMN mn [] t = ILam EmptyFC MW ExplicitArg (Just mn) (Implicit EmptyFC False) t

115 |     stringLamMN mn ns t = `(\ (Builtins.MkPair ~(IBindVar EmptyFC mn) ~(stringBind ns)) => ~t)

117 |     mkEither : (i,n : Nat) -> TTImp -> TTImp

118 |     mkEither Z (S (S _)) t = `(Prelude.Left ~t)

119 |     mkEither (S n) (S n') t = `(Prelude.Right ~(mkEither n n' t))

120 |     mkEither _ _ t = t

122 |     arrowImplLam : SnocList TTImp -> List ArrowVar -> TTImp -> Elab (SnocList TTImp)

123 |     arrowImpl' : SnocList TTImp -> List ArrowVar -> TTImp -> Elab (SnocList TTImp)

125 |     arrowImplLam ts [] (ILam _ MW ExplicitArg Nothing _ rest) =

126 |       arrowImpl' (ts :< `(Control.Arrow.arrow {a=_,b=_} (\_ => Builtin.MkUnit))) [] rest

127 |     arrowImplLam ts [] (ILam _ MW ExplicitArg (Just $ UN $ Basic var) ty rest) =

128 |       arrowImpl' ts [MkAVar var ty] rest

129 |     arrowImplLam ts strings (ILam _ MW ExplicitArg Nothing _ rest) =

130 |       arrowImpl' (ts :< `(Control.Arrow.arrow {a=_,b=_} Builtin.snd)) strings rest

131 |     arrowImplLam ts strings (ILam _ MW ExplicitArg (Just $ UN $ Basic var) ty rest) =

132 |       if any (\v => v.name == var) strings

133 |       then do

134 |         let used = map (\v => v.name /= var) strings

135 |         let strings' = MkAVar var ty :: zipFilter strings used

136 |         arrowImpl'

137 |           (ts :< `(Control.Arrow.arrow {a = ~(pairTy $ MkAVar var ty :: strings), b = ~(pairTy strings')}

138 |                     (\(Builtin.MkPair ~(IBindVar EmptyFC $ UN $ Basic var) ~(stringBind $ zip strings used)) =>

139 |                       ~(pair $ map name strings'))))

140 |           strings' rest

141 |       else arrowImpl' ts (MkAVar var ty :: strings) rest

142 |     arrowImplLam ts strings (ILam _ MW ExplicitArg (Just $ MN var _) ty rest) = do

143 |       arrowImpl' ts (MkAVar var ty :: strings) rest

144 |     arrowImplLam ts strings (ILam fc {}) = failAt fc "Invalid lambda"

145 |     arrowImplLam ts strings t = failAt (getFC t) "Expected lambda expression or binding pattern"

147 |     arrowImpl' ts (MkAVar v ty :: strings) (ICase fc _ (IVar _ var@(MN {})) _ clauses) = do

148 |       let tot = length clauses

149 |       (clauses', conts) <- map unzip $ evalStateT Z $ for clauses $ \case

150 |         PatClause fc' lhs rhs => do

151 |           let names = bindVars lhs

152 |           let vars = map (`MkAVar` `(_)) names

153 |           let usedLater = map (\v => not (elem v.name names) && usedInDiagram rhs v.name) strings

154 |           let strings' = vars ++ zipFilter strings usedLater

155 |           t <- lift $ assert_total $ arrowImpl' [<] strings' rhs

156 |           i <- get

157 |           modify S

158 |           pure (PatClause fc' lhs $ mkEither i tot $ pair $ map name strings', Just $ composeImp $ optimize t)

159 |         ImpossibleClause fc' lhs => pure (ImpossibleClause fc' lhs, Nothing)

160 |         WithClause fc' {} => failAt fc' "Unrecognized case pattern"

161 |       let conts' = catMaybes conts

162 |       case conts' of

163 |         [] =>

164 |           pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy $ MkAVar v ty :: strings), b = _}

165 |                       ~(stringLamMN var (map (,False) strings)

166 |                         (ICase fc [] (IVar EmptyFC var) ty clauses'))))

167 |         _ :: _ =>

168 |           pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy $ MkAVar v ty :: strings), b = _}

169 |                       ~(stringLamMN var (map (,False) strings)

170 |                         (ICase fc [] (IVar EmptyFC var) ty clauses')))

171 |                     :< foldr1 (\t,t' => `(Control.Arrow.(\|/) ~t ~t')) conts')

172 |     arrowImpl' ts strings (ICase fc _ exp ty clauses) = do

173 |       let tot = length clauses

174 |       (clauses', conts) <- map unzip $ evalStateT Z $ for clauses $ \case

175 |         PatClause fc' lhs rhs => do

176 |           let names = bindVars lhs

177 |           let vars = map (`MkAVar` `(_)) names

178 |           let usedLater = map (\v => not (elem v.name names) && usedInDiagram rhs v.name) strings

179 |           let strings' = vars ++ zipFilter strings usedLater

180 |           t <- lift $ assert_total $ arrowImpl' [<] strings' rhs

181 |           i <- get

182 |           modify S

183 |           pure (PatClause fc' lhs $ mkEither i tot $ pair $ map name strings', Just $ composeImp $ optimize t)

184 |         ImpossibleClause fc' lhs => pure (ImpossibleClause fc' lhs, Nothing)

185 |         WithClause fc' {} => failAt fc' "Unrecognized case pattern"

186 |       let conts' = catMaybes conts

187 |       case conts' of

188 |         [] =>

189 |           pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

190 |                       ~(stringLam (map (,True) strings) (ICase fc [] exp ty clauses'))))

191 |         _ :: _ =>

192 |           pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

193 |                       ~(stringLam (map (,True) strings) (ICase fc [] exp ty clauses')))

194 |                     :< foldr1 (\t,t' => `(Control.Arrow.(\|/) ~t ~t')) conts')

195 |     arrowImpl' ts strings `(~(ICase fc _ exp ty clauses) >> ~(rest)) = do

196 |       let tot = length clauses

197 |       (clauses', conts) <- map unzip $ evalStateT Z $ for clauses $ \case

198 |         PatClause fc' lhs rhs => do

199 |           let names = bindVars lhs

200 |           let vars = map (`MkAVar` `(_)) names

201 |           let usedLater = map (\v => not (elem v.name names) && usedInDiagram rhs v.name) strings

202 |           let strings' = vars ++ zipFilter strings usedLater

203 |           t <- lift $ assert_total $ arrowImpl' [<] strings' rhs

204 |           i <- get

205 |           modify S

206 |           pure (PatClause fc' lhs $ mkEither i tot $ pair $ map name strings', Just $ composeImp $ optimize t)

207 |         ImpossibleClause fc' lhs => pure (ImpossibleClause fc' lhs, Nothing)

208 |         WithClause fc' {} => failAt fc' "Unrecognized case pattern"

209 |       let conts' = catMaybes conts

210 |       case conts' of

211 |         [] =>

212 |           arrowImpl' (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

213 |                       ~(stringLam (map (,True) strings) (ICase fc [] exp ty clauses'))))

214 |             strings rest

215 |         _ :: _ =>

216 |           arrowImpl' (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

217 |                         ~(stringLam (map (,True) strings)

218 |                         `(Builtin.MkPair ~(ICase fc [] exp ty clauses') ~(pair $ map name strings))))

219 |                       :< `(Control.Arrow.first ~(foldr1 (\t,t' => `(Control.Arrow.(\|/) ~t ~t')) conts'))

220 |                       :< `(Control.Arrow.arrow {a=_,b=_} Builtin.snd))

221 |             strings rest

222 |     arrowImpl' ts strings `(~(ICase fc _ exp ty clauses) >>= ~(rest)) = do

223 |       let tot = length clauses

224 |       (clauses', conts) <- map unzip $ evalStateT Z $ for clauses $ \case

225 |         PatClause fc' lhs rhs => do

226 |           let names = bindVars lhs

227 |           let vars = map (`MkAVar` `(_)) names

228 |           let usedLater = map (\v => not (elem v.name names) && usedInDiagram rhs v.name) strings

229 |           let strings' = vars ++ zipFilter strings usedLater

230 |           t <- lift $ assert_total $ arrowImpl' [<] strings' rhs

231 |           i <- get

232 |           modify S

233 |           pure (PatClause fc' lhs $ mkEither i tot $ pair $ map name strings', Just $ composeImp $ optimize t)

234 |         ImpossibleClause fc' lhs => pure (ImpossibleClause fc' lhs, Nothing)

235 |         WithClause fc' {} => failAt fc' "Unrecognized case pattern"

236 |       let conts' = catMaybes conts

237 |       case conts' of

238 |         [] =>

239 |           arrowImplLam (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

240 |                       ~(stringLam (map (,True) strings) (ICase fc [] exp ty clauses'))))

241 |             strings rest

242 |         _ :: _ =>

243 |           arrowImplLam (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

244 |                         ~(stringLam (map (,True) strings)

245 |                         `(Builtin.MkPair ~(ICase fc [] exp ty clauses') ~(pair $ map name strings))))

246 |                       :< `(Control.Arrow.first ~(foldr1 (\t,t' => `(Control.Arrow.(\|/) ~t ~t')) conts')))

247 |             strings rest

248 |     arrowImpl' ts strings (ILet fc fc' MW (UN $ Basic var) ty exp rest) = do

249 |       (usedExp,usedLater,strings') <- do

250 |         let usedExp = map (usedInExpr exp . name) strings

251 |         let usedLater = map (\v => v.name /= var && usedInDiagram rest v.name) strings

252 |         let strings' = MkAVar var ty :: zipFilter strings usedLater

253 |         pure (usedExp,usedLater,strings')

254 |       arrowImpl'

255 |         (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = ~(pairTy strings')}

256 |           ~(stringLam (zip strings $ zipWith (\x,y => x || y) usedExp usedLater)

257 |             (ILet fc fc' MW (UN $ Basic var) ty exp (pair $ map name strings')))))

258 |         strings' rest

259 |     arrowImpl' ts strings (ILet fc {}) = failAt fc "Let binding must be unrestricted"

260 |     arrowImpl' ts strings `((~(mor) -< ~(inp)) >>= ~(rest)) =

261 |       case strings of

262 |         [] => arrowImplLam (ts :< `(Control.Arrow.arrow {a=_,b=_} (\_ => ~inp)) :< mor) [] rest

263 |         _ :: _ => do

264 |           (usedInp,usedLater,strings') <- do

265 |             let usedInp = map (usedInExpr inp . name) strings

266 |             let usedLater = map (usedInDiagram rest . name) strings

267 |             let strings' = zipFilter strings usedLater

268 |             pure (usedInp,usedLater,strings')

269 |           case strings' of

270 |             [] =>

271 |               arrowImplLam

272 |                 (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings),b=_}

273 |                           ~(stringLam (zip strings usedInp) inp))

274 |                     :< mor)

275 |                 strings' rest

276 |             _ :: _ =>

277 |               arrowImplLam

278 |                 (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = Builtin.Pair _ ~(pairTy strings')}

279 |                   ~(stringLam (zip strings $ zipWith (\x,y => x || y) usedInp usedLater)

280 |                     `(Builtin.MkPair ~inp ~(pair $ map name strings'))))

281 |                     :< `(Control.Arrow.first ~mor))

282 |                 strings' rest

283 |     arrowImpl' ts strings `((~(mor) -< ~(inp)) >> ~(rest)) =

284 |       case strings of

285 |         [] => arrowImpl' (ts

286 |           :< `(Control.Arrow.arrow {a=Builtin.Unit,b=_} (\_ => ~inp))

287 |           :< mor

288 |           :< `(Control.Arrow.arrow {a=_,b=Builtin.Unit} (\_ => Builtin.MkUnit))) [] rest

289 |         _ :: _ => do

290 |           (usedInp,usedLater,strings') <- do

291 |             let usedInp = map (usedInExpr inp . name) strings

292 |             let usedLater = map (usedInDiagram rest . name) strings

293 |             let strings' = zipFilter strings usedLater

294 |             pure (usedInp,usedLater,strings')

295 |           case strings' of

296 |             [] =>

297 |               arrowImpl'

298 |                 (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

299 |                           ~(stringLam (zip strings usedInp) inp))

300 |                     :< mor :< `(Control.Arrow.arrow {a=_,b=_} (\_ => ())))

301 |                 strings' rest

302 |             _ :: _ =>

303 |               arrowImpl'

304 |                 (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = Builtin.Pair _ ~(pairTy strings')}

305 |                         ~(stringLam (zip strings $ zipWith (\x,y => x || y) usedInp usedLater)

306 |                           `(Builtin.MkPair ~inp ~(pair $ map name strings'))))

307 |                     :< `(Control.Arrow.first ~mor) :< `(Control.Arrow.arrow {a=_,b=_} Builtin.snd))

308 |                 strings' rest

309 |     arrowImpl' ts strings `(~(mor) -< ~(inp)) =

310 |       pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings), b = _}

311 |                     ~(stringLam (map (,True) strings) inp))

312 |                 :< mor)

313 |     arrowImpl' ts strings `(=< ~(out)) =

314 |       pure (ts :< `(Control.Arrow.arrow {a = ~(pairTy strings),b=_}

315 |         ~(stringLam (map (,True) strings) out)))

317 |     arrowImpl' _ _ (IHole fc _) = failAt fc "Holes are not allowed as top-level commands"

318 |     arrowImpl' ts strings (ILam fc {}) = failAt (getFC t) "Lambda not allowed here"

319 |     arrowImpl' ts strings t = failAt (getFC t) "Could not parse expression"

326 | arrowDo : {0 arr : Hom Type0} -> (0 _ : Arrow arr) => TTImp -> Elab (arr (W0 a) (W0 b))

327 | arrowDo t = check !(arrowImpl t)