IN: inference
USING: generic interpreter kernel lists math namespaces
sequences words ;

: infer-shuffle-inputs ( shuffle node -- )
    >r dup shuffle-in-d length swap shuffle-in-r length r>
    node-inputs ;

: shuffle-stacks ( shuffle -- )
    #! Shuffle simulated stacks.
    meta-d get meta-r get rot shuffle meta-r set meta-d set ;

: infer-shuffle-outputs ( shuffle node -- )
    >r dup shuffle-out-d length swap shuffle-out-r length r>
    node-outputs ;

: infer-shuffle ( shuffle -- )
    #shuffle
    2dup infer-shuffle-inputs
    over shuffle-stacks
    tuck infer-shuffle-outputs
    node, ;

: shuffle>effect ( shuffle -- effect )
    dup shuffle-in-d [ drop object ] map
    swap shuffle-out-d [ drop object ] map 2list ;

: define-shuffle ( word shuffle -- )
    [ shuffle>effect "infer-effect" set-word-prop ] 2keep
    [ , \ infer-shuffle , ] [ ] make "infer" set-word-prop ;

@{
    @{ drop  << shuffle f 1 0 @{             }@ @{   }@ >> }@
    @{ 2drop << shuffle f 2 0 @{             }@ @{   }@ >> }@
    @{ 3drop << shuffle f 3 0 @{             }@ @{   }@ >> }@
    @{ dup   << shuffle f 1 0 @{ 0 0         }@ @{   }@ >> }@
    @{ 2dup  << shuffle f 2 0 @{ 0 1 0 1     }@ @{   }@ >> }@
    @{ 3dup  << shuffle f 3 0 @{ 0 1 2 0 1 2 }@ @{   }@ >> }@
    @{ rot   << shuffle f 3 0 @{ 1 2 0       }@ @{   }@ >> }@
    @{ -rot  << shuffle f 3 0 @{ 2 0 1       }@ @{   }@ >> }@
    @{ dupd  << shuffle f 2 0 @{ 0 0 1       }@ @{   }@ >> }@
    @{ swapd << shuffle f 3 0 @{ 1 0 2       }@ @{   }@ >> }@
    @{ nip   << shuffle f 2 0 @{ 1           }@ @{   }@ >> }@
    @{ 2nip  << shuffle f 3 0 @{ 2           }@ @{   }@ >> }@
    @{ tuck  << shuffle f 2 0 @{ 1 0 1       }@ @{   }@ >> }@
    @{ over  << shuffle f 2 0 @{ 0 1 0       }@ @{   }@ >> }@
    @{ pick  << shuffle f 3 0 @{ 0 1 2 0     }@ @{   }@ >> }@
    @{ swap  << shuffle f 2 0 @{ 1 0         }@ @{   }@ >> }@
    @{ >r    << shuffle f 1 0 @{             }@ @{ 0 }@ >> }@
    @{ r>    << shuffle f 0 1 @{ 0           }@ @{   }@ >> }@
}@ [ first2 define-shuffle ] each