HaskellImport.hs

-- | Deterministic imports and external references for generated Haskell.
--
-- A renderer supplies every external reference needed by one target module.
-- This module makes all qualification choices over that complete set so the
-- resulting aliases and bytes do not depend on declaration traversal order.
module Keiro.Dsl.HaskellImport
  ( HaskellNamespace (..),
    QualificationPreference (..),
    HaskellReference (..),
    ImportEnvironment (..),
    HaskellImportError (..),
    HaskellImportPlan,
    planHaskellImports,
    renderPlannedImports,
    renderPlannedReference,
  )
where

import Data.Foldable (traverse_)
import Data.List (find)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Keiro.Dsl.HaskellName (haskellKeywords)

data HaskellNamespace
  = TypeNamespace
  | ValueNamespace
  | ConstructorNamespace
  deriving stock (Eq, Ord, Show)

data QualificationPreference
  = PreferUnqualified
  | RequireQualified
  deriving stock (Eq, Ord, Show)

data HaskellReference = HaskellReference
  { moduleName :: !Text,
    name :: !Text,
    namespace :: !HaskellNamespace,
    qualification :: !QualificationPreference
  }
  deriving stock (Eq, Ord, Show)

data ImportEnvironment = ImportEnvironment
  { targetModule :: !Text,
    localNames :: !(Set Text),
    reservedQualifiers :: !(Set Text)
  }
  deriving stock (Eq, Show)

data HaskellImportError
  = InvalidHaskellModule !Text !Text
  | InvalidHaskellOccurrence !Text !HaskellNamespace !Text
  | HaskellSelfImport !Text
  | ImpossibleHaskellAlias !Text !(Set Text)
  | MissingHaskellReference !Text !HaskellReference
  deriving stock (Eq, Show)

data HaskellImportPlan = HaskellImportPlan
  { targetModule :: !Text,
    declarations :: !(Set Text),
    references :: !(Map HaskellReference Text)
  }

planHaskellImports :: ImportEnvironment -> Set HaskellReference -> Either HaskellImportError HaskellImportPlan
planHaskellImports environment references = do
  validateModuleName target target
  traverse_ (validateReference target) orderedReferences
  aliases <- allocateAliases environment unqualifiedNames qualifiedModules
  plannedReferences <- Map.fromList <$> traverse (renderReference aliases) orderedReferences
  pure
    HaskellImportPlan
      { targetModule = target,
        declarations = explicitDeclarations <> qualifiedDeclarations aliases,
        references = plannedReferences
      }
  where
    target = (.targetModule) environment
    orderedReferences = Set.toAscList references
    occurrenceOwners =
      Map.fromListWith
        (<>)
        [ ((.name) reference, Set.singleton ((.moduleName) reference, (.name) reference))
        | reference <- orderedReferences,
          (.namespace) reference == TypeNamespace,
          (.qualification) reference == PreferUnqualified
        ]
    unqualified reference =
      (.namespace) reference == TypeNamespace
        && (.qualification) reference == PreferUnqualified
        && Set.notMember ((.name) reference) ((.localNames) environment)
        && Set.notMember ((.name) reference) ((.reservedQualifiers) environment)
        && maybe False ((== 1) . Set.size) (Map.lookup ((.name) reference) occurrenceOwners)
    unqualifiedReferences = Set.filter unqualified references
    unqualifiedNames = Set.map (.name) unqualifiedReferences
    qualifiedModules =
      Set.map (.moduleName) (references `Set.difference` unqualifiedReferences)
    explicitImports =
      Map.fromListWith
        (<>)
        [ ((.moduleName) reference, Set.singleton ((.name) reference))
        | reference <- Set.toAscList unqualifiedReferences
        ]
    explicitDeclarations =
      Set.fromList
        [ "import " <> moduleName <> " (" <> T.intercalate ", " (Set.toAscList names) <> ")"
        | (moduleName, names) <- Map.toAscList explicitImports
        ]
    qualifiedDeclarations aliases =
      Set.fromList
        [ "import " <> moduleName <> " qualified as " <> alias
        | (moduleName, alias) <- Map.toAscList aliases
        ]
    renderReference aliases reference
      | Set.member reference unqualifiedReferences = pure (reference, (.name) reference)
      | otherwise = case Map.lookup ((.moduleName) reference) aliases of
          Nothing -> Left (ImpossibleHaskellAlias target (Set.singleton ((.moduleName) reference)))
          Just alias -> pure (reference, alias <> "." <> (.name) reference)

renderPlannedImports :: HaskellImportPlan -> Text
renderPlannedImports = T.intercalate "\n" . Set.toAscList . (.declarations)

renderPlannedReference :: HaskellImportPlan -> HaskellReference -> Either HaskellImportError Text
renderPlannedReference plan reference =
  maybe
    (Left (MissingHaskellReference ((.targetModule) plan) reference))
    Right
    (Map.lookup reference ((.references) plan))

allocateAliases :: ImportEnvironment -> Set Text -> Set Text -> Either HaskellImportError (Map Text Text)
allocateAliases environment unqualifiedNames modules = do
  let moduleCandidates = Map.fromSet suffixCandidates modules
      candidateOwners =
        Map.fromListWith
          (<>)
          [ (candidate, Set.singleton moduleName)
          | (moduleName, candidates) <- Map.toAscList moduleCandidates,
            candidate <- candidates
          ]
      occupied = (.reservedQualifiers) environment <> (.localNames) environment <> unqualifiedNames
      choose moduleName candidates =
        case find (isAvailable moduleName candidateOwners occupied) candidates of
          Just candidate -> candidate
          Nothing -> T.intercalate "_" (moduleComponents moduleName)
      aliases = Map.mapWithKey choose moduleCandidates
      ownersByAlias =
        Map.fromListWith
          (<>)
          [ (alias, Set.singleton moduleName)
          | (moduleName, alias) <- Map.toAscList aliases
          ]
      impossibleModules =
        Set.unions
          [ owners
          | (alias, owners) <- Map.toAscList ownersByAlias,
            Set.member alias occupied || Set.size owners /= 1 || not (validUpperIdentifier alias)
          ]
  if Set.null impossibleModules
    then pure aliases
    else Left (ImpossibleHaskellAlias ((.targetModule) environment) impossibleModules)

isAvailable :: Text -> Map Text (Set Text) -> Set Text -> Text -> Bool
isAvailable moduleName candidateOwners occupied candidate =
  Set.notMember candidate occupied
    && Map.lookup candidate candidateOwners == Just (Set.singleton moduleName)
    && validUpperIdentifier candidate

suffixCandidates :: Text -> [Text]
suffixCandidates moduleName =
  [ T.concat (drop (componentCount - suffixLength) components)
  | suffixLength <- [1 .. componentCount]
  ]
  where
    components = moduleComponents moduleName
    componentCount = length components

moduleComponents :: Text -> [Text]
moduleComponents = T.splitOn "."

validateReference :: Text -> HaskellReference -> Either HaskellImportError ()
validateReference target reference = do
  validateModuleName target ((.moduleName) reference)
  if (.moduleName) reference == target
    then Left (HaskellSelfImport target)
    else pure ()
  if validOccurrence ((.namespace) reference) ((.name) reference)
    then pure ()
    else Left (InvalidHaskellOccurrence target ((.namespace) reference) ((.name) reference))

validateModuleName :: Text -> Text -> Either HaskellImportError ()
validateModuleName target candidate
  | not (null components) && all validUpperIdentifier components = pure ()
  | otherwise = Left (InvalidHaskellModule target candidate)
  where
    components = moduleComponents candidate

validOccurrence :: HaskellNamespace -> Text -> Bool
validOccurrence namespace name =
  not (Set.member name haskellKeywords)
    && case namespace of
      TypeNamespace -> validUpperIdentifier name
      ConstructorNamespace -> validUpperIdentifier name
      ValueNamespace -> validLowerIdentifier name

validUpperIdentifier :: Text -> Bool
validUpperIdentifier name = case T.uncons name of
  Just (first, rest) -> asciiUpper first && T.all identifierTail rest
  Nothing -> False

validLowerIdentifier :: Text -> Bool
validLowerIdentifier name = case T.uncons name of
  Just (first, rest) -> (asciiLower first || first == '_') && T.all identifierTail rest
  Nothing -> False

identifierTail :: Char -> Bool
identifierTail character =
  asciiUpper character
    || asciiLower character
    || asciiDigit character
    || character == '_'
    || character == '\''

asciiUpper :: Char -> Bool
asciiUpper character = character >= 'A' && character <= 'Z'

asciiLower :: Char -> Bool
asciiLower character = character >= 'a' && character <= 'z'

asciiDigit :: Char -> Bool
asciiDigit character = character >= '0' && character <= '9'