{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DerivingStrategies #-}
module Glean.Query.BindOrder
  ( Fix
  , Scope(..)
  , NoBind(..)
  , IsPat(..)
  , runFixBindOrder
  , FixBindOrder(..)
  , FixBindOrderError(..)
  , Variable
  , isBound
  , bind
  ) where

import Control.Monad.Except
import Control.Monad.State
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap

import Glean.Query.Codegen.Types (Match(..), Var(..), Generator_(..), Generator)
import Glean.Query.Vars (VarSet)
import Glean.RTS.Term as RTS

-- -----------------------------------------------------------------------------
-- Fixing up MatchBind vs. MatchVar

-- Substitution can mess up MatchBind and MatchVar - we might end up
-- with multiple MatchBinds for a variable, or a MatchVar before a
-- MatchBind. To fix it up all we need to do is traverse the query in
-- the correct order, keeping track of which variables are in scope,
-- and change MatchBind<->MatchVar as appropriate.
--
-- This is all somewhat suboptimal, because the typechecker has
-- already figured out MatchBind vs. MatchVar and here we mess it up
-- by substitution and then fix it again. Which begs the question: why
-- do we have MatchBind and MatchVar at all, couldn't we leave it
-- until the last minute just before code generation to figure out
-- which variables are binding occurrences?  Yes, but it's nice to be
-- able to give the user out-of-scope error messages from the
-- typechecker.  Maybe we'll change this in the future.

data Scope = Scope
  { isScope :: VarSet
  , bound :: IntMap Var
  }
  deriving (Show)

type Variable = Int

isBound :: Scope -> Variable -> Bool
isBound (Scope _ bound) var = IntMap.member var bound

bind :: Var -> Scope -> Scope
bind var (Scope scope bound) = Scope scope (IntMap.insert (varId var) var bound)


-- | Variables that should not be bound.
-- Statements in a negated subquery should not bind values that occur later
-- in a positive statement.
newtype NoBind = NoBind
  { unNoBind :: IntSet
  }

data FixBindOrderError
  = CannotUseWildcardInExpr
  | CannotUseNeverInExpr
  | UnboundVariable Var

type Fix a = StateT (Scope, NoBind) (Except FixBindOrderError) a

runFixBindOrder
  :: Scope -> NoBind -> Fix a -> Except FixBindOrderError (a, Scope)
runFixBindOrder scope nobind fx = do
  (res, (scope', _)) <- runStateT fx (scope, nobind)
  return (res, scope')

data IsPat = IsPat | IsExpr
  deriving Eq

class FixBindOrder a where
  fixBindOrder :: IsPat -> a -> Fix a

instance FixBindOrder Generator where
  fixBindOrder isPat (TermGenerator p) =
    TermGenerator <$> fixBindOrder isPat p
  fixBindOrder _ (ArrayElementGenerator ty p) =
    ArrayElementGenerator ty <$> fixBindOrder IsExpr p
  fixBindOrder _ (SetElementGenerator ty p) =
    SetElementGenerator ty <$> fixBindOrder IsExpr p
  fixBindOrder isPat (DerivedFactGenerator pid key val) =
    DerivedFactGenerator pid
      <$> fixBindOrder isPat key
      <*> fixBindOrder isPat val
  fixBindOrder _ (FactGenerator pid kpat vpat range) =
    FactGenerator pid <$>
      fixBindOrder IsPat kpat <*>
      fixBindOrder IsPat vpat <*>
      pure range
  fixBindOrder _ (PrimCall op pats ty) =
    PrimCall op <$> mapM (fixBindOrder IsExpr) pats <*> pure ty

instance FixBindOrder a => FixBindOrder (Term a) where
  fixBindOrder isPat term = mapM (fixBindOrder isPat) term

instance FixBindOrder () where
  fixBindOrder _ () = return ()

instance (FixBindOrder a) => FixBindOrder (Match a Var) where
  fixBindOrder IsExpr MatchWild{} =
    throwError CannotUseWildcardInExpr
  fixBindOrder IsExpr MatchNever{} =
    throwError CannotUseNeverInExpr
  fixBindOrder isPat (MatchBind var) = fixVar isPat var
  fixBindOrder isPat (MatchVar var) = fixVar isPat var
  fixBindOrder isPat (MatchAnd a b) =
    MatchAnd <$> mapM (fixBindOrder isPat) a <*> mapM (fixBindOrder isPat) b
  fixBindOrder isPat (MatchPrefix str rest) =
    MatchPrefix str <$> mapM (fixBindOrder isPat) rest
  fixBindOrder isPat (MatchArrayPrefix ty pre all) =
    MatchArrayPrefix ty
      <$> (mapM.mapM) (fixBindOrder isPat) pre
      <*> mapM (fixBindOrder isPat) all
  fixBindOrder isPat (MatchExt ext) = MatchExt <$> fixBindOrder isPat ext
  fixBindOrder _ other@MatchFid{} = return other
  fixBindOrder _ other@MatchNever{} = return other
  fixBindOrder _ other@MatchWild{} = return other

fixVar :: IsPat -> Var -> Fix (Match a Var)
fixVar isPat var@(Var _ v _) = do
  (scope, NoBind noBind) <- get
  if
    | isBound scope v -> return (MatchVar var)
    | (IsExpr == isPat) || (v `IntSet.member` noBind)
      -> throwError $ UnboundVariable var
    | otherwise -> do
      put (bind var scope, NoBind noBind)
      return (MatchBind var)