Paste: infix operations

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 
prettyprint ;
IN: infix

<PRIVATE
: 2compose ( quot quot quot -- compose ) compose compose ; inline

: prepare-operand ( term -- quot )
    dup callable? [ '[ _ call ] ] [ 1quotation ] if ;

: construct-op ( l r quot -- result/quot )
    2over [ number? ] both? [ call ] [ 
        [ [ prepare-operand ] bi@ ] dip 2compose
    ] 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-zA-Z_0-9]* "[" expression "]" =>
    [[ [ fourth prepare-operand ] [ first2 swap prefix >local-word 1quotation ] bi [ nth ] 2compose ]]

fun = [a-zA-Z] [a-zA-Z_0-9]* "(" (funargs)? ")" =>
    [[ [ fourth dup . [ [ ] ] unless* ] keep first2 swap prefix >string find-and-check ]]

funargs = (funargs ",")? expression =>
    [[ [ third prepare-operand ] [ first ] bi [ compose ] when* ]]

ident = [a-zA-z] [a-zA-Z_0-9]* => [[ first2 swap prefix >local-word ]]

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

sum-term = (sum-term ("-"|"+"))? product-term =>
    [[ first2 over [ [ first2 ] dip swap "+" = [ +op ] [ -op ] if ] [ nip ] if ]]

product-term = (product-term ("*"|"/"|"%"))? (paren-exp|terminal) =>
    [[ first2 over [ [ first2 ] dip swap select-product-op ] [ nip ] if ]]

paren-exp = "(" expression ")" => [[ second ]]
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