module Unison.Typechecker.Components (minimize, minimize') where

import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NEL
import Data.Map qualified as Map
import Data.Set qualified as Set
import Data.Set.NonEmpty (NESet)
import Data.Set.NonEmpty qualified as NESet
import Unison.ABT qualified as ABT
import Unison.Prelude
import Unison.Term (Term')
import Unison.Term qualified as Term
import Unison.Var (Var)
import Unison.Var qualified as Var

unordered :: (Var v) => [(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
unordered :: forall v vt a.
Var v =>
[(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
unordered = [(v, Term (F vt a a) v a)] -> [[(v, Term (F vt a a) v a)]]
forall v (f :: * -> *) a.
Var v =>
[(v, Term f v a)] -> [[(v, Term f v a)]]
ABT.components

ordered :: (Var v) => [(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
ordered :: forall v vt a.
Var v =>
[(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
ordered = [(v, Term (F vt a a) v a)] -> [[(v, Term (F vt a a) v a)]]
forall v (f :: * -> *) a.
Var v =>
[(v, Term f v a)] -> [[(v, Term f v a)]]
ABT.orderedComponents

-- | Algorithm for minimizing cycles of a `let rec`. This can
-- improve generalization during typechecking and may also be more
-- efficient for execution.
--
-- For instance:
--
-- minimize (let rec id x = x; g = id 42; y = id "hi" in g)
-- ==>
-- Just (let id x = x; g = id 42; y = id "hi" in g)
--
-- Gets rid of the let rec and replaces it with an ordinary `let`, such
-- that `id` is suitably generalized.
--
-- Fails on the left if there are duplicate definitions.
minimize ::
  forall vt v a.
  (Var v, Ord a) =>
  Term' vt v a ->
  Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
minimize :: forall vt v a.
(Var v, Ord a) =>
Term' vt v a
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
minimize (Term.LetRecNamedAnnotatedTop' Bool
isTop a
blockAnn [((a, v), Term' vt v a)]
bs Term' vt v a
e) =
  let bindings :: [(v, Term' vt v a)]
      bindings :: [(v, Term' vt v a)]
bindings = [((a, v), Term' vt v a)]
bs [((a, v), Term' vt v a)]
-> (((a, v), Term' vt v a) -> (v, Term' vt v a))
-> [(v, Term' vt v a)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \((a
_a, v
v), Term' vt v a
t) -> (v
v, Term' vt v a
t)
      grouped :: Map v (NESet a)
      grouped :: Map v (NESet a)
grouped =
        [((a, v), Term' vt v a)]
bs
          [((a, v), Term' vt v a)]
-> ([((a, v), Term' vt v a)] -> [(v, NESet a)]) -> [(v, NESet a)]
forall a b. a -> (a -> b) -> b
& (((a, v), Term' vt v a) -> (v, NESet a))
-> [((a, v), Term' vt v a)] -> [(v, NESet a)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\((a
a, v
v), Term' vt v a
_t) -> (v
v, a -> NESet a
forall a. a -> NESet a
NESet.singleton a
a)) -- For duplicates, we care about the binding location.
          [(v, NESet a)]
-> ([(v, NESet a)] -> Map v (NESet a)) -> Map v (NESet a)
forall a b. a -> (a -> b) -> b
& (NESet a -> NESet a -> NESet a)
-> [(v, NESet a)] -> Map v (NESet a)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith (NESet a -> NESet a -> NESet a
forall a. Ord a => NESet a -> NESet a -> NESet a
NESet.union)
      dupes :: Map v (NESet a)
dupes = (v -> NESet a -> Bool) -> Map v (NESet a) -> Map v (NESet a)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey v -> NESet a -> Bool
forall {v} {t :: * -> *} {a}.
(Var v, Foldable t) =>
v -> t a -> Bool
ok Map v (NESet a)
grouped
        where
          ok :: v -> t a -> Bool
ok v
v t a
as
            | v -> Text
forall v. Var v => v -> Text
Var.name v
v Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"_" = Bool
False
            | Bool
otherwise = t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
as Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
   in case [(v, NESet a)] -> Maybe (NonEmpty (v, NESet a))
forall a. [a] -> Maybe (NonEmpty a)
NEL.nonEmpty ([(v, NESet a)] -> Maybe (NonEmpty (v, NESet a)))
-> [(v, NESet a)] -> Maybe (NonEmpty (v, NESet a))
forall a b. (a -> b) -> a -> b
$ Map v (NESet a) -> [(v, NESet a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map v (NESet a)
dupes of
        Just NonEmpty (v, NESet a)
dupeList -> NonEmpty (v, NESet a)
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall a b. a -> Either a b
Left NonEmpty (v, NESet a)
dupeList
        Maybe (NonEmpty (v, NESet a))
Nothing -> do
          let cs0 :: [[(v, Term' vt v a)]]
cs0 = if Bool
isTop then [(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
forall v vt a.
Var v =>
[(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
unordered [(v, Term' vt v a)]
bindings else [(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
forall v vt a.
Var v =>
[(v, Term' vt v a)] -> [[(v, Term' vt v a)]]
ordered [(v, Term' vt v a)]
bindings
              -- within a cycle, we put the lambdas first, so
              -- unguarded definitions can refer to these lambdas, example:
              --
              --   foo x = blah + 1 + x
              --   blah = foo 10
              --
              -- Here `foo` and `blah` are part of a cycle, but putting `foo`
              -- first at least lets the program run (though it has an infinite
              -- loop).
              cs :: [[(v, Term' vt v a)]]
cs = ((v, Term' vt v a) -> Bool)
-> [(v, Term' vt v a)] -> [(v, Term' vt v a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (\(v
_, Term' vt v a
e) -> Term' vt v a -> Int
forall vt at ap v a. Term2 vt at ap v a -> Int
Term.arity Term' vt v a
e Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) ([(v, Term' vt v a)] -> [(v, Term' vt v a)])
-> [[(v, Term' vt v a)]] -> [[(v, Term' vt v a)]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[(v, Term' vt v a)]]
cs0
              varAnnotations :: Map v a
varAnnotations = [(v, a)] -> Map v a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((\((a
a, v
v), Term' vt v a
_) -> (v
v, a
a)) (((a, v), Term' vt v a) -> (v, a))
-> [((a, v), Term' vt v a)] -> [(v, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [((a, v), Term' vt v a)]
bs)
              msg :: v -> a
msg v
v = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ [Char]
"Components.minimize " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> (v, [v]) -> [Char]
forall a. Show a => a -> [Char]
show (v
v, Map v a -> [v]
forall k a. Map k a -> [k]
Map.keys Map v a
varAnnotations)
              annotationFor :: v -> a
annotationFor v
v = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe (v -> a
msg v
v) (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$ v -> Map v a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup v
v Map v a
varAnnotations
              annotatedVar :: v -> (a, v)
annotatedVar v
v = (v -> a
annotationFor v
v, v
v)
              -- When introducing a nested let/let rec, we use the annotation
              -- of the variable that starts off that let/let rec
              mklet :: [(v, Term' vt v a)] -> Term' vt v a -> Term' vt v a
mklet [(v
hdv, Term' vt v a
hdb)] Term' vt v a
e
                | v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
hdv (Term' vt v a -> Set v
forall (f :: * -> *) v a. Term f v a -> Set v
ABT.freeVars Term' vt v a
hdb) =
                    Bool
-> a -> [((a, v), Term' vt v a)] -> Term' vt v a -> Term' vt v a
forall v vt a.
Ord v =>
Bool
-> a -> [((a, v), Term' vt v a)] -> Term' vt v a -> Term' vt v a
Term.letRec
                      Bool
isTop
                      a
blockAnn
                      [(v -> (a, v)
annotatedVar v
hdv, Term' vt v a
hdb)]
                      Term' vt v a
e
                | Bool
otherwise = Bool -> a -> a -> (v, Term' vt v a) -> Term' vt v a -> Term' vt v a
forall v a vt at ap.
Ord v =>
Bool
-> a
-> a
-> (v, Term2 vt at ap v a)
-> Term2 vt at ap v a
-> Term2 vt at ap v a
Term.singleLet Bool
isTop a
blockAnn (v -> a
annotationFor v
hdv) (v
hdv, Term' vt v a
hdb) Term' vt v a
e
              mklet cycle :: [(v, Term' vt v a)]
cycle@((v
_, Term' vt v a
_) : [(v, Term' vt v a)]
_) Term' vt v a
e =
                Bool
-> a -> [((a, v), Term' vt v a)] -> Term' vt v a -> Term' vt v a
forall v vt a.
Ord v =>
Bool
-> a -> [((a, v), Term' vt v a)] -> Term' vt v a -> Term' vt v a
Term.letRec
                  Bool
isTop
                  a
blockAnn
                  ((v -> (a, v)) -> (v, Term' vt v a) -> ((a, v), Term' vt v a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first v -> (a, v)
annotatedVar ((v, Term' vt v a) -> ((a, v), Term' vt v a))
-> [(v, Term' vt v a)] -> [((a, v), Term' vt v a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(v, Term' vt v a)]
cycle)
                  Term' vt v a
e
              mklet [] Term' vt v a
e = Term' vt v a
e
           in Maybe (Term' vt v a)
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall a b. b -> Either a b
Right (Maybe (Term' vt v a)
 -> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a)))
-> ([[(v, Term' vt v a)]] -> Maybe (Term' vt v a))
-> [[(v, Term' vt v a)]]
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term' vt v a -> Maybe (Term' vt v a)
forall a. a -> Maybe a
Just (Term' vt v a -> Maybe (Term' vt v a))
-> ([[(v, Term' vt v a)]] -> Term' vt v a)
-> [[(v, Term' vt v a)]]
-> Maybe (Term' vt v a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(v, Term' vt v a)] -> Term' vt v a -> Term' vt v a)
-> Term' vt v a -> [[(v, Term' vt v a)]] -> Term' vt v a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr [(v, Term' vt v a)] -> Term' vt v a -> Term' vt v a
mklet Term' vt v a
e ([[(v, Term' vt v a)]]
 -> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a)))
-> [[(v, Term' vt v a)]]
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall a b. (a -> b) -> a -> b
$ [[(v, Term' vt v a)]]
cs
minimize Term' vt v a
_ = Maybe (Term' vt v a)
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall a b. b -> Either a b
Right Maybe (Term' vt v a)
forall a. Maybe a
Nothing

minimize' ::
  (Var v, Ord a) => Term' vt v a -> Either (NonEmpty (v, NESet a)) (Term' vt v a)
minimize' :: forall v a vt.
(Var v, Ord a) =>
Term' vt v a -> Either (NonEmpty (v, NESet a)) (Term' vt v a)
minimize' Term' vt v a
term = Term' vt v a -> Maybe (Term' vt v a) -> Term' vt v a
forall a. a -> Maybe a -> a
fromMaybe Term' vt v a
term (Maybe (Term' vt v a) -> Term' vt v a)
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
-> Either (NonEmpty (v, NESet a)) (Term' vt v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term' vt v a
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
forall vt v a.
(Var v, Ord a) =>
Term' vt v a
-> Either (NonEmpty (v, NESet a)) (Maybe (Term' vt v a))
minimize Term' vt v a
term