    001/*
    002 * Copyright (c) 2007-2017 Xplenty, Inc. All Rights Reserved.
    003 *
    004 * Project and contact information: http://www.cascading.org/
    005 *
    006 * This file is part of the Cascading project.
    007 *
    008 * Licensed under the Apache License, Version 2.0 (the "License");
    009 * you may not use this file except in compliance with the License.
    010 * You may obtain a copy of the License at
    011 *
    012 *     http://www.apache.org/licenses/LICENSE-2.0
    013 *
    014 * Unless required by applicable law or agreed to in writing, software
    015 * distributed under the License is distributed on an "AS IS" BASIS,
    016 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    017 * See the License for the specific language governing permissions and
    018 * limitations under the License.
    019 */
    020
    021package cascading.flow.planner.iso.transformer;
    022
    023import java.util.Set;
    024
    025import cascading.flow.FlowElement;
    026import cascading.flow.planner.graph.AnnotatedGraph;
    027import cascading.flow.planner.graph.ElementGraph;
    028import cascading.flow.planner.graph.ElementMultiGraph;
    029import cascading.flow.planner.iso.ElementAnnotation;
    030import cascading.flow.planner.iso.expression.ExpressionGraph;
    031import cascading.flow.planner.iso.finder.Match;
    032import cascading.util.ProcessLogger;
    033
    034/**
    035 *
    036 */
    037public class AnnotateGraphTransformer extends RecursiveGraphTransformer<ElementGraph>
    038  {
    039  private final ElementAnnotation annotation;
    040
    041  protected final GraphTransformer graphTransformer;
    042
    043  public AnnotateGraphTransformer( ExpressionGraph match, ElementAnnotation annotation )
    044    {
    045    super( match );
    046    this.annotation = annotation;
    047    this.graphTransformer = null;
    048    }
    049
    050  public AnnotateGraphTransformer( GraphTransformer graphTransformer, ExpressionGraph match, ElementAnnotation annotation )
    051    {
    052    super( match );
    053    this.annotation = annotation;
    054    this.graphTransformer = graphTransformer;
    055    }
    056
    057  @Override
    058  protected boolean requiresRecursiveSearch()
    059    {
    060    return graphTransformer != null || super.requiresRecursiveSearch();
    061    }
    062
    063  @Override
    064  protected ElementGraph prepareForMatch( ProcessLogger processLogger, Transformed<ElementGraph> transformed, ElementGraph graph )
    065    {
    066    if( graphTransformer == null )
    067      return makeAnnotated( graph );
    068
    069    Transformed child = graphTransformer.transform( transformed.getPlannerContext(), graph );
    070
    071    transformed.addChildTransform( child );
    072
    073    ElementGraph endGraph = child.getEndGraph();
    074
    075    return makeAnnotated( endGraph );
    076    }
    077
    078  private ElementGraph makeAnnotated( ElementGraph endGraph )
    079    {
    080    if( endGraph == null )
    081      return null;
    082
    083    if( endGraph instanceof AnnotatedGraph )
    084      return endGraph;
    085
    086    return new ElementMultiGraph( endGraph );
    087    }
    088
    089  @Override
    090  protected Set<FlowElement> addExclusions( ElementGraph graph )
    091    {
    092    return ( (AnnotatedGraph) graph ).getAnnotations().getValues( annotation.getAnnotation() );
    093    }
    094
    095  @Override
    096  protected boolean transformGraphInPlaceUsing( Transformed<ElementGraph> transformed, ElementGraph graph, Match match )
    097    {
    098    Set<FlowElement> captured = match.getCapturedElements( annotation.getCapture() );
    099
    100    if( captured.isEmpty() )
    101      return false;
    102
    103    ( (AnnotatedGraph) graph ).getAnnotations().addAll( annotation.getAnnotation(), captured );
    104
    105    return true;
    106    }
    107
    108  }