Jruby behavior nodes

package org.jruby.truffle.core.behavior;


import com.oracle.truffle.api.CompilerDirectives;
import com.oracle.truffle.api.dsl.Specialization;
import com.oracle.truffle.api.frame.VirtualFrame;
import com.oracle.truffle.api.nodes.Node;
import com.oracle.truffle.api.source.SourceSection;
import org.jruby.truffle.builtins.CoreClass;
import org.jruby.truffle.builtins.CoreMethod;
import org.jruby.truffle.builtins.CoreMethodNode;
import org.jruby.truffle.builtins.UnaryCoreMethodNode;
import org.jruby.truffle.core.*;
import org.jruby.truffle.core.behavior.init.InitBehaviorExpression;
import org.jruby.truffle.core.behavior.init.InitFold;
import org.jruby.truffle.core.behavior.utility.BehaviorOption;
import org.jruby.truffle.core.behavior.utility.DependencyStaticScope;
import org.jruby.truffle.language.NotProvided;
import org.jruby.truffle.language.dispatch.CallDispatchHeadNode;
import org.jruby.truffle.language.dispatch.DispatchHeadNodeFactory;
import org.jruby.truffle.language.objects.WriteObjectFieldNode;
import org.jruby.truffle.RubyContext;
import org.jruby.RubyProc;
import org.jruby.truffle.BehaviorObject;
import org.jruby.truffle.language.objects.WriteObjectFieldNodeGen;

@CoreClass("BehaviorCore")
public class BehaviorModule {

    @CoreMethod(names = "fold", isModuleFunction = true, required = 1, needsBlock = true  )
    public abstract static class FoldExprNode extends CoreMethodNode {

        @Node.Child
        DependencyStaticScope extractDeps;
        @Node.Child
        InitFold initFold;

        public FoldExprNode(RubyContext context, SourceSection sourceSection) {
            super(context, sourceSection);
            extractDeps = new DependencyStaticScope();
            initFold = new InitFold(context);
        }

        @Override
        public Object execute(VirtualFrame frame) {
            return NotProvided.INSTANCE;
        }

        @Specialization
        public BehaviorObject fold(VirtualFrame frame, int value, RubyProc proc){
            BehaviorObject[] deps = extractDeps.execute(proc);
            return initFold.execute(frame,deps,value,proc);
        }
        @Specialization
        public BehaviorObject fold(VirtualFrame frame, double value, RubyProc proc){
            BehaviorObject[] deps = extractDeps.execute(proc);
            return initFold.execute(frame,deps,value,proc);        }
        @Specialization
        public BehaviorObject fold(VirtualFrame frame, Object value, RubyProc proc){
            BehaviorObject[] deps = extractDeps.execute(proc);
            return initFold.execute(frame,deps,value,proc);
        }
    }

    @CoreMethod(names = {"behavior","signal","map"}, isModuleFunction = true, needsBlock = true)
    public abstract static class BehaviorExprNode extends UnaryCoreMethodNode {
        @Child
        private WriteObjectFieldNode writeSignalExpr;
        @Child
        InitBehaviorExpression execSignalExpr;
        @Child
        DependencyStaticScope extractDeps;

        public BehaviorExprNode(RubyContext context, SourceSection sourceSection) {
            super(context, sourceSection);
            writeSignalExpr = WriteObjectFieldNodeGen.create(BehaviorOption.SIGNAL_EXPR);
            execSignalExpr = new InitBehaviorExpression(context, sourceSection);
            extractDeps = new DependencyStaticScope();
        }

        @Specialization
        BehaviorObject map(VirtualFrame frame, RubyProc block) {
            BehaviorObject self = newSignal();
            BehaviorObject[] dependsOn = extractDeps.execute(block);
            self.setupPropagationDep(dependsOn);
            writeSignalExpr.execute(self, block);
            execSignalExpr.execute(frame, self, dependsOn);
            return self;
        }

        @CompilerDirectives.TruffleBoundary
        private BehaviorObject newSignal() {
            return new BehaviorObject(getContext());
            /*return (BehaviorObject) //FIXME correct?
                    (new BehaviorObject.SignalRuntimeAllocator()).allocate(
                            getContext(),
                            getContext().getCoreLibrary().getBehaviorClass(),
                            null);*/
        }
    }


    @CoreMethod(names = "source", isModuleFunction = true, required = 1)
    public abstract static class SourceNode extends UnaryCoreMethodNode {
        @Child
        CallDispatchHeadNode callInit;


        public SourceNode(RubyContext context, SourceSection sourceSection) {
            super(context, sourceSection);
            callInit = DispatchHeadNodeFactory.createMethodCall(context, true);
        }

        @Specialization
        BehaviorObject source(VirtualFrame frame, int value) {
            final BehaviorObject self = newSignal();
            return (BehaviorObject) callInit.call(frame, newSignal(), "initialize", null, value);
        }

        @Specialization
        BehaviorObject source(VirtualFrame frame, Object value) {
            final BehaviorObject self = newSignal();
            return (BehaviorObject) callInit.call(frame, newSignal(), "initialize", null, value);
        }

        @CompilerDirectives.TruffleBoundary
        private BehaviorObject newSignal() {
            return new BehaviorObject(getContext());
            /* FIXME correct?
            return (BehaviorObject) (new BehaviorObject.SignalRuntimeAllocator()).allocate(
                    getContext(),
                    getContext().getCoreLibrary().getBehaviorSourceClass(),
                    null);*/
        }

    }

}