module Unison.Syntax.FilePrinter
  ( renderDefnsForUnisonFile,
  )
where

import Control.Lens (mapped, _1)
import Control.Monad.Writer (Writer)
import Control.Monad.Writer qualified as Writer
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Unison.Builtin.Decls qualified as Builtin.Decls
import Unison.ConstructorReference (ConstructorReference, GConstructorReference (..))
import Unison.DataDeclaration (Decl)
import Unison.DeclNameLookup (DeclNameLookup, expectConstructorNames)
import Unison.HashQualified qualified as HQ
import Unison.HashQualifiedPrime qualified as HQ'
import Unison.Name (Name)
import Unison.Prelude
import Unison.PrettyPrintEnvDecl (PrettyPrintEnvDecl (..))
import Unison.Reference (TypeReferenceId)
import Unison.Reference qualified as Reference
import Unison.Referent (Referent)
import Unison.Referent qualified as Referent
import Unison.Syntax.DeclPrinter (AccessorName)
import Unison.Syntax.DeclPrinter qualified as DeclPrinter
import Unison.Syntax.TermPrinter qualified as TermPrinter
import Unison.Term (Term)
import Unison.Type (Type)
import Unison.Typechecker qualified as Typechecker
import Unison.Util.Defns (Defns (..), DefnsF)
import Unison.Util.Pretty (ColorText, Pretty)
import Unison.Util.Pretty qualified as Pretty
import Unison.Var (Var)

-- | Render definitions destined for a Unison file.
--
-- This first renders the types (discovering which record accessors will be generated upon parsing), then renders the
-- terms (being careful not to render any record accessors, since those would cause duplicate binding errors upon
-- parsing).
renderDefnsForUnisonFile ::
  forall a v.
  (Var v, Monoid a) =>
  DeclNameLookup ->
  PrettyPrintEnvDecl ->
  DefnsF (Map Name) (Term v a, Type v a) (TypeReferenceId, Decl v a) ->
  DefnsF (Map Name) (Pretty ColorText) (Pretty ColorText)
renderDefnsForUnisonFile :: forall a v.
(Var v, Monoid a) =>
DeclNameLookup
-> PrettyPrintEnvDecl
-> DefnsF
     (Map Name) (Term v a, Type v a) (TypeReferenceId, Decl v a)
-> DefnsF (Map Name) (Pretty ColorText) (Pretty ColorText)
renderDefnsForUnisonFile DeclNameLookup
declNameLookup PrettyPrintEnvDecl
ppe DefnsF (Map Name) (Term v a, Type v a) (TypeReferenceId, Decl v a)
defns =
  let (Map Name (Pretty ColorText)
types, Set Name
accessorNames) = Writer (Set Name) (Map Name (Pretty ColorText))
-> (Map Name (Pretty ColorText), Set Name)
forall w a. Writer w a -> (a, w)
Writer.runWriter ((Name
 -> (TypeReferenceId, Decl v a)
 -> WriterT (Set Name) Identity (Pretty ColorText))
-> Map Name (TypeReferenceId, Decl v a)
-> Writer (Set Name) (Map Name (Pretty ColorText))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey Name
-> (TypeReferenceId, Decl v a)
-> WriterT (Set Name) Identity (Pretty ColorText)
renderType DefnsF (Map Name) (Term v a, Type v a) (TypeReferenceId, Decl v a)
defns.types)
   in Defns
        { $sel:terms:Defns :: Map Name (Pretty ColorText)
terms = (Name -> (Term v a, Type v a) -> Maybe (Pretty ColorText))
-> Map Name (Term v a, Type v a) -> Map Name (Pretty ColorText)
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybeWithKey (Set Name
-> Name -> (Term v a, Type v a) -> Maybe (Pretty ColorText)
renderTerm Set Name
accessorNames) DefnsF (Map Name) (Term v a, Type v a) (TypeReferenceId, Decl v a)
defns.terms,
          Map Name (Pretty ColorText)
types :: Map Name (Pretty ColorText)
$sel:types:Defns :: Map Name (Pretty ColorText)
types
        }
  where
    renderType :: Name -> (TypeReferenceId, Decl v a) -> Writer (Set AccessorName) (Pretty ColorText)
    renderType :: Name
-> (TypeReferenceId, Decl v a)
-> WriterT (Set Name) Identity (Pretty ColorText)
renderType Name
name (TypeReferenceId
ref, Decl v a
typ) =
      (Pretty SyntaxText -> Pretty ColorText)
-> WriterT (Set Name) Identity (Pretty SyntaxText)
-> WriterT (Set Name) Identity (Pretty ColorText)
forall a b.
(a -> b)
-> WriterT (Set Name) Identity a -> WriterT (Set Name) Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Pretty SyntaxText -> Pretty ColorText
forall r. Pretty (SyntaxText' r) -> Pretty ColorText
Pretty.syntaxToColor (WriterT (Set Name) Identity (Pretty SyntaxText)
 -> WriterT (Set Name) Identity (Pretty ColorText))
-> WriterT (Set Name) Identity (Pretty SyntaxText)
-> WriterT (Set Name) Identity (Pretty ColorText)
forall a b. (a -> b) -> a -> b
$
        PrettyPrintEnvDecl
-> Reference' Text Hash
-> HashQualified Name
-> Decl v a
-> WriterT (Set Name) Identity (Pretty SyntaxText)
forall v a.
Var v =>
PrettyPrintEnvDecl
-> Reference' Text Hash
-> HashQualified Name
-> Decl v a
-> WriterT (Set Name) Identity (Pretty SyntaxText)
DeclPrinter.prettyDeclW
          -- Sort of a hack; since the decl printer looks in the PPE for names of constructors,
          -- we just delete all term names out and add back the constructors...
          -- probably no need to wipe out the suffixified side but we do it anyway
          (DeclNameLookup
-> Name
-> TypeReferenceId
-> PrettyPrintEnvDecl
-> PrettyPrintEnvDecl
setPpedToConstructorNames DeclNameLookup
declNameLookup Name
name TypeReferenceId
ref PrettyPrintEnvDecl
ppe)
          (TypeReferenceId -> Reference' Text Hash
Reference.fromId TypeReferenceId
ref)
          (Name -> HashQualified Name
forall n. n -> HashQualified n
HQ.NameOnly Name
name)
          Decl v a
typ

    renderTerm :: Set Name -> Name -> (Term v a, Type v a) -> Maybe (Pretty ColorText)
    renderTerm :: Set Name
-> Name -> (Term v a, Type v a) -> Maybe (Pretty ColorText)
renderTerm Set Name
accessorNames Name
name (Term v a
term, Type v a
typ) = do
      Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Bool
not (Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Name
name Set Name
accessorNames))
      let hqName :: HashQualified Name
hqName = Name -> HashQualified Name
forall n. n -> HashQualified n
HQ.NameOnly Name
name
      let rendered :: Pretty SyntaxText
rendered
            | Type v a -> Type v a -> Bool
forall v loc. Var v => Type v loc -> Type v loc -> Bool
Typechecker.isEqual (a -> Type v a
forall v a. Ord v => a -> Type v a
Builtin.Decls.testResultListType a
forall a. Monoid a => a
mempty) Type v a
typ =
                Pretty SyntaxText
"test> " Pretty SyntaxText -> Pretty SyntaxText -> Pretty SyntaxText
forall a. Semigroup a => a -> a -> a
<> PrettyPrintEnv
-> HashQualified Name -> Term v a -> Pretty SyntaxText
forall v at ap a.
Var v =>
PrettyPrintEnv
-> HashQualified Name -> Term2 v at ap v a -> Pretty SyntaxText
TermPrinter.prettyBindingWithoutTypeSignature PrettyPrintEnvDecl
ppe.suffixifiedPPE HashQualified Name
hqName Term v a
term
            | Bool
otherwise = PrettyPrintEnv
-> HashQualified Name -> Term v a -> Pretty SyntaxText
forall v at ap a.
Var v =>
PrettyPrintEnv
-> HashQualified Name -> Term2 v at ap v a -> Pretty SyntaxText
TermPrinter.prettyBinding PrettyPrintEnvDecl
ppe.suffixifiedPPE HashQualified Name
hqName Term v a
term
      Pretty ColorText -> Maybe (Pretty ColorText)
forall a. a -> Maybe a
Just (Pretty SyntaxText -> Pretty ColorText
forall r. Pretty (SyntaxText' r) -> Pretty ColorText
Pretty.syntaxToColor Pretty SyntaxText
rendered)

setPpedToConstructorNames :: DeclNameLookup -> Name -> TypeReferenceId -> PrettyPrintEnvDecl -> PrettyPrintEnvDecl
setPpedToConstructorNames :: DeclNameLookup
-> Name
-> TypeReferenceId
-> PrettyPrintEnvDecl
-> PrettyPrintEnvDecl
setPpedToConstructorNames DeclNameLookup
declNameLookup Name
name TypeReferenceId
ref =
  ASetter
  PrettyPrintEnvDecl
  PrettyPrintEnvDecl
  (Referent -> [(HashQualified Name, HashQualified Name)])
  (Referent -> [(HashQualified Name, HashQualified Name)])
-> (Referent -> [(HashQualified Name, HashQualified Name)])
-> PrettyPrintEnvDecl
-> PrettyPrintEnvDecl
forall s t a b. ASetter s t a b -> b -> s -> t
set ((PrettyPrintEnv -> Identity PrettyPrintEnv)
-> PrettyPrintEnvDecl -> Identity PrettyPrintEnvDecl
#unsuffixifiedPPE ((PrettyPrintEnv -> Identity PrettyPrintEnv)
 -> PrettyPrintEnvDecl -> Identity PrettyPrintEnvDecl)
-> (((Referent -> [(HashQualified Name, HashQualified Name)])
     -> Identity
          (Referent -> [(HashQualified Name, HashQualified Name)]))
    -> PrettyPrintEnv -> Identity PrettyPrintEnv)
-> ASetter
     PrettyPrintEnvDecl
     PrettyPrintEnvDecl
     (Referent -> [(HashQualified Name, HashQualified Name)])
     (Referent -> [(HashQualified Name, HashQualified Name)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Referent -> [(HashQualified Name, HashQualified Name)])
 -> Identity
      (Referent -> [(HashQualified Name, HashQualified Name)]))
-> PrettyPrintEnv -> Identity PrettyPrintEnv
#termNames) Referent -> [(HashQualified Name, HashQualified Name)]
referentNames
    (PrettyPrintEnvDecl -> PrettyPrintEnvDecl)
-> (PrettyPrintEnvDecl -> PrettyPrintEnvDecl)
-> PrettyPrintEnvDecl
-> PrettyPrintEnvDecl
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter
  PrettyPrintEnvDecl
  PrettyPrintEnvDecl
  (Referent -> [(HashQualified Name, HashQualified Name)])
  (Referent -> [(HashQualified Name, HashQualified Name)])
-> (Referent -> [(HashQualified Name, HashQualified Name)])
-> PrettyPrintEnvDecl
-> PrettyPrintEnvDecl
forall s t a b. ASetter s t a b -> b -> s -> t
set ((PrettyPrintEnv -> Identity PrettyPrintEnv)
-> PrettyPrintEnvDecl -> Identity PrettyPrintEnvDecl
#suffixifiedPPE ((PrettyPrintEnv -> Identity PrettyPrintEnv)
 -> PrettyPrintEnvDecl -> Identity PrettyPrintEnvDecl)
-> (((Referent -> [(HashQualified Name, HashQualified Name)])
     -> Identity
          (Referent -> [(HashQualified Name, HashQualified Name)]))
    -> PrettyPrintEnv -> Identity PrettyPrintEnv)
-> ASetter
     PrettyPrintEnvDecl
     PrettyPrintEnvDecl
     (Referent -> [(HashQualified Name, HashQualified Name)])
     (Referent -> [(HashQualified Name, HashQualified Name)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Referent -> [(HashQualified Name, HashQualified Name)])
 -> Identity
      (Referent -> [(HashQualified Name, HashQualified Name)]))
-> PrettyPrintEnv -> Identity PrettyPrintEnv
#termNames) Referent -> [(HashQualified Name, HashQualified Name)]
referentNames
  where
    constructorNameMap :: Map ConstructorReference Name
    constructorNameMap :: Map ConstructorReference Name
constructorNameMap =
      [(ConstructorReference, Name)] -> Map ConstructorReference Name
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
        ( Name
name
            Name -> (Name -> [Name]) -> [Name]
forall a b. a -> (a -> b) -> b
& HasCallStack => DeclNameLookup -> Name -> [Name]
DeclNameLookup -> Name -> [Name]
expectConstructorNames DeclNameLookup
declNameLookup
            [Name] -> ([Name] -> [(Word64, Name)]) -> [(Word64, Name)]
forall a b. a -> (a -> b) -> b
& [Word64] -> [Name] -> [(Word64, Name)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip [Word64
0 ..]
            [(Word64, Name)]
-> ([(Word64, Name)] -> [(ConstructorReference, Name)])
-> [(ConstructorReference, Name)]
forall a b. a -> (a -> b) -> b
& ASetter
  [(Word64, Name)]
  [(ConstructorReference, Name)]
  Word64
  ConstructorReference
-> (Word64 -> ConstructorReference)
-> [(Word64, Name)]
-> [(ConstructorReference, Name)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (((Word64, Name) -> Identity (ConstructorReference, Name))
-> [(Word64, Name)] -> Identity [(ConstructorReference, Name)]
Setter
  [(Word64, Name)]
  [(ConstructorReference, Name)]
  (Word64, Name)
  (ConstructorReference, Name)
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped (((Word64, Name) -> Identity (ConstructorReference, Name))
 -> [(Word64, Name)] -> Identity [(ConstructorReference, Name)])
-> ((Word64 -> Identity ConstructorReference)
    -> (Word64, Name) -> Identity (ConstructorReference, Name))
-> ASetter
     [(Word64, Name)]
     [(ConstructorReference, Name)]
     Word64
     ConstructorReference
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word64 -> Identity ConstructorReference)
-> (Word64, Name) -> Identity (ConstructorReference, Name)
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (Word64, Name)
  (ConstructorReference, Name)
  Word64
  ConstructorReference
_1) (Reference' Text Hash -> Word64 -> ConstructorReference
forall r. r -> Word64 -> GConstructorReference r
ConstructorReference (TypeReferenceId -> Reference' Text Hash
Reference.fromId TypeReferenceId
ref))
        )

    referentNames :: Referent -> [(HQ'.HashQualified Name, HQ'.HashQualified Name)]
    referentNames :: Referent -> [(HashQualified Name, HashQualified Name)]
referentNames = \case
      Referent.Con ConstructorReference
conRef ConstructorType
_ ->
        case ConstructorReference -> Map ConstructorReference Name -> Maybe Name
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ConstructorReference
conRef Map ConstructorReference Name
constructorNameMap of
          Maybe Name
Nothing -> []
          Just Name
conName -> let hqConName :: HashQualified Name
hqConName = Name -> HashQualified Name
forall n. n -> HashQualified n
HQ'.NameOnly Name
conName in [(HashQualified Name
hqConName, HashQualified Name
hqConName)]
      Referent.Ref Reference' Text Hash
_ -> []