USING: ascii assocs accessors combinators combinators.short-circuit
effects fry kernel math math.parser multiline namespaces parser
peg.ebnf sequences strings words locals.parser locals.types
quotations stack-checker summary ;
IN: infix

<PRIVATE
: prepare-operand ( term -- quot )
    dup callable? [ 1quotation ] unless ;

: construct-op ( l r quot -- result/quot )
    2over [ number? ] both? [ call ] [ 
        [ [ prepare-operand ] bi@ ] dip '[ @ @ @ ]
    ] if ; inline

: +op ( l r -- result/quot ) [ + ] construct-op ; inline
: -op ( l r -- result/quot ) [ - ] construct-op ; inline
: *op ( l r -- result/quot ) [ * ] construct-op ; inline
: /op ( l r -- result/quot ) [ / ] construct-op ; inline
: %op ( l r -- result/quot ) [ mod ] construct-op ; inline
PRIVATE>

ERROR: local-not-defined name ;
M: local-not-defined summary
    drop "local is not defined" ;
ERROR: word-not-found name ;
M: word-not-found summary
    drop "Word not found in the current vocabulary search pathaoe" ;
ERROR: bad-stack-effect ;

<PRIVATE
: at? ( key assoc -- value/key ? ) dupd at* [ nip t ] [ drop f ] if ;

: >local-word ( seq -- word )
    >string locals get at? [ local-not-defined ] unless ;

: select-sum-op ( l r opstring -- quot )
    "+" = [ +op ] [ -op ] if ; inline

: select-product-op ( l r opstring -- quot )
    {
        { "*" [ *op ] }
        { "/" [ /op ] }
        [ drop %op ]
    } case ; inline

: check-word ( args word -- ? )
    {
        [ nip [ optimized>> ] [ primitive? ] bi or ]
        [
            [ infer out>> ]
            [ stack-effect [ in>> length ] [ out>> length ] bi ] bi*
            [ = ] [ 1 = ] bi* and
        ] 
    } 2&& ;

: find-and-check ( args string -- quot )
    dup search [ nip ] [ no-word ] if*
    2dup check-word [ bad-stack-effect ] unless
    1quotation compose ;

EBNF: infix-parser
arr = [a-zA-Z]:a [a-zA-Z_0-9]*:b "[" expression:e "]" =>
    [[ e prepare-operand b a prefix >local-word 1quotation '[ @ @ nth ] ]]

fun = [a-zA-Z]:a [a-zA-Z_0-9]*:b "(" (funargs)?:arguments ")" =>
    [[ arguments [ [ ] ] unless* b a prefix >string find-and-check ]]

funargs = (funargs ",")?:args expression:e =>
    [[ e prepare-operand args [ first compose ] when* ]]

ident = [a-zA-z]:a [a-zA-Z_0-9]*:b => [[ b a prefix >local-word ]]

num = ("-"|"+")? [0-9]+ "."? [0-9]* => [[ concat >string string>number ]]

sum-term = (sum-term ("-"|"+"))?:leftpart product-term:r =>
    [[ leftpart [ first2 r swap select-sum-op ] [ r ] if* ]]

product-term = (product-term ("*"|"/"|"%"))?:leftpart (paren-exp|terminal):r =>
    [[ leftpart [ first2 r swap select-product-op ] [ r ] if* ]]

paren-exp = "(" expression:e ")" => [[ e ]]
terminal = num | arr | fun | ident
expression = sum-term | terminal
;EBNF

: [infix-parse ( end -- result/quot )
    parse-multiline-string [ blank? not ] filter infix-parser
    dup number? [ 1quotation ] when ;
PRIVATE>

DEFER: infix]
: [infix
    "infix]" [infix-parse parsed \ call parsed ; parsing

<PRIVATE
: parse-infix-locals ( assoc end -- quot )
    [
        in-lambda? on
        [ dup [ locals set ] [ push-locals ] bi ] dip
        [infix-parse [ call ] curry swap pop-locals
    ] with-scope ;
PRIVATE>

: [infix|
    "|" parse-bindings "infix]" parse-infix-locals <let>
    parsed-lambda ; parsing