ParserUtils.hs

{-# OPTIONS  #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Language.Python.Common.ParserUtils 
-- Copyright   : (c) 2009 Bernie Pope 
-- License     : BSD-style
-- Maintainer  : bjpop@csse.unimelb.edu.au
-- Stability   : experimental
-- Portability : ghc
--
-- Various utilities to support the Python parser. 
-----------------------------------------------------------------------------

module Language.Python.Common.ParserUtils where

import Data.List (foldl')
import Data.Maybe (isJust)
import Language.Python.Common.AST as AST
import Language.Python.Common.Token as Token 
import Language.Python.Common.ParserMonad hiding (location)
import Language.Python.Common.SrcLocation 

makeConditionalExpr :: ExprSpan -> Maybe (ExprSpan, ExprSpan) -> ExprSpan
makeConditionalExpr e Nothing = e
makeConditionalExpr e opt@(Just (cond, false_branch))
   = CondExpr e cond false_branch (spanning e opt)

makeBinOp :: ExprSpan -> [(OpSpan, ExprSpan)] -> ExprSpan
makeBinOp e es
   = foldl' mkOp e es
   where
   mkOp e1 (op, e2) = BinaryOp op e1 e2 (spanning e1 e2)

parseError :: Token -> P a 
parseError = throwError . UnexpectedToken 

data Trailer
   = TrailerCall { trailer_call_args :: [ArgumentSpan], trailer_span :: SrcSpan }
   | TrailerSubscript { trailer_subs :: [Subscript], trailer_span :: SrcSpan }
   | TrailerDot { trailer_dot_ident :: IdentSpan, dot_span :: SrcSpan, trailer_span :: SrcSpan }

instance Span Trailer where
  getSpan = trailer_span

data Subscript
   = SubscriptExpr { subscription :: ExprSpan, subscript_span :: SrcSpan }
   | SubscriptSlice 
     { subscript_slice_span1 :: Maybe ExprSpan
     , subscript_slice_span2 :: Maybe ExprSpan
     , subscript_slice_span3 :: Maybe (Maybe ExprSpan)
     , subscript_span :: SrcSpan
     }
   | SubscriptSliceEllipsis { subscript_span :: SrcSpan }

instance Span Subscript where
   getSpan = subscript_span

isProperSlice :: Subscript -> Bool
isProperSlice (SubscriptSlice {}) = True
isProperSlice (SubscriptSliceEllipsis {}) = True
isProperSlice other = False

subscriptToSlice :: Subscript -> SliceSpan
subscriptToSlice (SubscriptSlice lower upper stride span)
   = SliceProper lower upper stride span
subscriptToSlice (SubscriptExpr e span)
   = SliceExpr e span
subscriptToSlice (SubscriptSliceEllipsis span)
   = SliceEllipsis span

subscriptToExpr :: Subscript -> ExprSpan
subscriptToExpr (SubscriptExpr { subscription = s }) = s
subscriptToExpr other = error "subscriptToExpr applied to non subscript"

subscriptsToExpr :: [Subscript] -> ExprSpan
subscriptsToExpr subs
   | length subs > 1 = Tuple (map subscriptToExpr subs) (getSpan subs)
   | length subs == 1 = subscriptToExpr $ head subs
   | otherwise = error "subscriptsToExpr: empty subscript list"

addTrailer :: ExprSpan -> [Trailer] -> ExprSpan
addTrailer
   = foldl' trail
   where
   trail :: ExprSpan -> Trailer -> ExprSpan
   -- XXX fix the span
   trail e trail@(TrailerCall { trailer_call_args = args }) = Call e args (spanning e trail)
   trail e trail@(TrailerSubscript { trailer_subs = subs })
      | any isProperSlice subs
           = SlicedExpr e (map subscriptToSlice subs) (spanning e trail) 
      | otherwise 
           = Subscript e (subscriptsToExpr subs) (spanning e trail) 
   trail e trail@(TrailerDot { trailer_dot_ident = ident, dot_span = ds })
      = Dot { dot_expr = e, dot_attribute = ident, expr_annot = spanning e trail }

makeTupleOrExpr :: [ExprSpan] -> Maybe Token -> ExprSpan
makeTupleOrExpr [e] Nothing = e
makeTupleOrExpr es@(_:_) (Just t) = Tuple es (spanning es t) 
makeTupleOrExpr es@(_:_) Nothing  = Tuple es (getSpan es)
makeTupleOrExpr [] _ = error "makeTupleOrExpr should never be called with an empty list"

makeAssignmentOrExpr :: ExprSpan -> Either [ExprSpan] (AssignOpSpan, ExprSpan) -> StatementSpan
makeAssignmentOrExpr e (Left es) 
   = makeNormalAssignment e es
makeAssignmentOrExpr e (Right ope2)
   = makeAugAssignment e ope2

makeAugAssignment :: ExprSpan -> (AssignOpSpan, ExprSpan) -> StatementSpan
makeAugAssignment e1 (op, e2)
  = AST.AugmentedAssign e1 op e2 (spanning e1 e2)

makeNormalAssignment :: ExprSpan -> [ExprSpan] -> StatementSpan
makeNormalAssignment e [] = StmtExpr e (getSpan e)
makeNormalAssignment e es
  = AST.Assign (e : front) (head back) (spanning e es)
  where
  (front, back) = splitAt (len - 1) es
  len = length es

makeAnnAssignment :: ExprSpan -> (ExprSpan, Maybe ExprSpan) -> StatementSpan
makeAnnAssignment ato (annotation, ae) = AST.AnnotatedAssign annotation ato ae (spanning ae ato)

makeTry :: Token -> SuiteSpan -> ([HandlerSpan], [StatementSpan], [StatementSpan]) -> StatementSpan
makeTry t1 body (handlers, elses, finally)
   = AST.Try body handlers elses finally 
     (spanning (spanning (spanning (spanning t1 body) handlers) elses) finally)

makeParam :: (IdentSpan, Maybe ExprSpan) -> Maybe ExprSpan -> ParameterSpan
makeParam (name, annot) defaultVal
   = Param name annot defaultVal paramSpan
   where
   paramSpan = spanning (spanning name annot) defaultVal

makeStarParam :: Token -> Maybe (IdentSpan, Maybe ExprSpan) -> ParameterSpan
makeStarParam t1 Nothing = EndPositional (getSpan t1) 
makeStarParam t1 (Just (name, annot))
   = VarArgsPos name annot (spanning t1 annot) 

makeStarStarParam :: Token -> (IdentSpan, Maybe ExprSpan) -> ParameterSpan
makeStarStarParam t1 (name, annot)
   = VarArgsKeyword name annot (spanning (spanning t1 name) annot) 

-- version 2 only 
makeTupleParam :: ParamTupleSpan -> Maybe ExprSpan -> ParameterSpan
-- just a name
makeTupleParam p@(ParamTupleName {}) optDefault = 
   Param (param_tuple_name p) Nothing optDefault (spanning p optDefault)
-- a parenthesised tuple. NOTE: we do not distinguish between (foo) and (foo,)
makeTupleParam p@(ParamTuple { param_tuple_annot = span }) optDefault =
   UnPackTuple p optDefault span 

makeComprehension :: ExprSpan -> CompForSpan -> ComprehensionSpan
makeComprehension e for = Comprehension (ComprehensionExpr e) for (spanning e for)

makeListForm :: SrcSpan -> Either ExprSpan ComprehensionSpan -> ExprSpan
makeListForm span (Left tuple@(Tuple {})) = List (tuple_exprs tuple) span
makeListForm span (Left other) = List [other] span 
makeListForm span (Right comprehension) = ListComp comprehension span

makeSet :: ExprSpan -> Either CompForSpan [ExprSpan] -> SrcSpan -> ExprSpan
makeSet e (Left compFor) = SetComp (Comprehension (ComprehensionExpr e) compFor (spanning e compFor))
makeSet e (Right es) = Set (e:es)

-- The Either (ExprSpan, ExprSpan) ExprSpan refers to a (key, value) pair or a dictionary unpacking expression.
makeDictionary :: Either (ExprSpan, ExprSpan) ExprSpan -> Either CompForSpan [Either (ExprSpan, ExprSpan) ExprSpan] -> SrcSpan -> ExprSpan
makeDictionary (Left mapping@(key, val)) (Left compFor) =
   DictComp (Comprehension (ComprehensionDict (DictMappingPair key val)) compFor (spanning mapping compFor))
-- This is allowed by the grammar, but will produce a runtime syntax error:
-- dict unpacking cannot be used in dict comprehension
makeDictionary (Right unpacking) (Left compFor) =
   DictComp (Comprehension (ComprehensionDict (DictUnpacking unpacking)) compFor (spanning unpacking compFor))
makeDictionary item (Right es) = Dictionary $ toKeyDatumList <$> item : es


toKeyDatumList :: Either (ExprSpan, ExprSpan) ExprSpan -> DictKeyDatumList SrcSpan
toKeyDatumList (Left (key, value)) = DictMappingPair key value
toKeyDatumList (Right unpacking) = DictUnpacking unpacking


fromEither :: Either a a -> a
fromEither (Left x) = x
fromEither (Right x) = x

makeDecorator :: Token -> DottedNameSpan -> [ArgumentSpan] -> DecoratorSpan
makeDecorator t1 name [] = Decorator name [] (spanning t1 name)
makeDecorator t1 name args = Decorator name args (spanning t1 args)

-- parser guarantees that the first list is non-empty
makeDecorated :: [DecoratorSpan] -> StatementSpan -> StatementSpan
makeDecorated ds@(d:_) def = Decorated ds def (spanning d def)
makeDecorated [] _ = error "parser guarantees that makeDecorated's first argument is non-empty"

-- suite can't be empty so it is safe to take span over it
makeFun :: Token -> IdentSpan -> [ParameterSpan] -> Maybe ExprSpan -> SuiteSpan -> StatementSpan
makeFun t1 name params annot body = 
   Fun name params annot body $ spanning t1 body 

makeReturn :: Token -> Maybe ExprSpan -> StatementSpan
makeReturn t1 Nothing = AST.Return Nothing (getSpan t1)
makeReturn t1 expr@(Just e) = AST.Return expr (spanning t1 e)

makeParenOrGenerator :: Either ExprSpan ComprehensionSpan -> SrcSpan -> ExprSpan
makeParenOrGenerator (Left e) span = Paren e span
makeParenOrGenerator (Right comp) span = Generator comp span

makePrint :: Bool -> Maybe ([ExprSpan], Maybe Token) -> SrcSpan -> StatementSpan
makePrint chevron Nothing span = AST.Print chevron [] False span
makePrint chevron (Just (args, last_comma)) span =
   AST.Print chevron args (isJust last_comma) span
   
makeRelative :: [Either Token DottedNameSpan] -> ImportRelativeSpan
makeRelative items =
   ImportRelative ndots maybeName (getSpan items) 
   where
   (ndots, maybeName) = countDots 0 items
   -- parser ensures that the dotted name will be at the end 
   -- of the list if it is there at all
   countDots :: Int -> [Either Token DottedNameSpan] -> (Int, Maybe DottedNameSpan)
   countDots count [] = (count, Nothing)
   countDots count (Right name:_) = (count, Just name)
   countDots count (Left token:rest) = countDots (count + dots token) rest 
   dots (DotToken {}) = 1
   dots (EllipsisToken {}) = 3
   dots _ = error "Parser ensures dots is only called on DotToken or EllipsisToken."

{-
   See: http://docs.python.org/3.0/reference/expressions.html#calls

   arglist: (argument ',')* (argument [',']
                         |'*' test (',' argument)* [',' '**' test]
                         |'**' test)

   (state 1) Positional arguments come first.
   (state 2) Then keyword arguments.
   (state 3) Then the single star form.
   (state 4) Then more keyword arguments (but no positional arguments).
   (state 5) Then the double star form.

XXX fixme: we need to include SrcLocations for the errors.
-}

checkArguments :: [ArgumentSpan] -> P [ArgumentSpan]
checkArguments args = do
   check 1 args
   return args
   where
   check :: Int -> [ArgumentSpan] -> P ()
   check state [] = return ()
   check 5 (arg:_) = spanError arg "an **argument must not be followed by any other arguments"
   check state (arg:rest) = do
      case arg of
         ArgExpr {}
            | state == 1 -> check state rest
            | state == 2 -> spanError arg "a positional argument must not follow a keyword argument"
            | otherwise -> spanError arg "a positional argument must not follow a *argument"
         ArgKeyword {}
            | state `elem` [1,2] -> check 2 rest
            | state `elem` [3,4] -> check 4 rest
            | otherwise -> error "state should always be in range 1..4 here"
         ArgVarArgsPos {}
            | state `elem` [1,2] -> check 3 rest
            | state `elem` [3,4] -> spanError arg "there must not be two *arguments in an argument list"
            | otherwise -> error "state should always be in range 1..4 here"
         ArgVarArgsKeyword {} -> check 5 rest

{-
   See: http://docs.python.org/3.1/reference/compound_stmts.html#grammar-token-parameter_list

   parameter_list ::=  (defparameter ",")*
                    (  "*" [parameter] ("," defparameter)*
                    [, "**" parameter]
                    | "**" parameter
                    | defparameter [","] )

   (state 1) Parameters/unpack tuples first.
   (state 2) Then the single star (on its own or with parameter)
   (state 3) Then more parameters. 
   (state 4) Then the double star form.

   XXX fixme, add support for version 2 unpack tuple.
-}

checkParameters :: [ParameterSpan] -> P [ParameterSpan]
checkParameters params = do
   check 1 params 
   return params
   where
   check :: Int -> [ParameterSpan] -> P ()
   check state [] = return ()
   check 4 (param:_) = spanError param "a **parameter must not be followed by any other parameters"
   check state (param:rest) = do
      case param of
         -- Param and UnPackTuple are treated the same.
         UnPackTuple {}
            | state `elem` [1,3] -> check state rest
            | state == 2 -> check 3 rest 
            | otherwise -> error "state should always be in range 1..3 here"
         Param {}
            | state `elem` [1,3] -> check state rest
            | state == 2 -> check 3 rest 
            | otherwise -> error "state should always be in range 1..3 here"
         EndPositional {}
            | state == 1 -> check 2 rest
            | otherwise -> spanError param "there must not be two *parameters in a parameter list"
         VarArgsPos {}
            | state == 1 -> check 2 rest
            | otherwise -> spanError param "there must not be two *parameters in a parameter list"
         VarArgsKeyword {} -> check 4 rest

{-
spanError :: Span a => a -> String -> P ()
spanError x str = throwError $ StrError $ unwords [prettyText $ getSpan x, str]
-}